{
    "cells": [
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "# Autoencoder Tic Tac Toe Verifier\n",
                "\n",
                "This is another approach of verifying games but using an anomaly detection approach instead of classification.\n",
                "\n",
                "Make sure to use a GPU otherwise you'll take a super long time to train"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "# check if notebook is in colab\n",
                "try:\n",
                "    import google.colab\n",
                "    import subprocess\n",
                "    import sys\n",
                "    subprocess.check_call([sys.executable, \"-m\", \"pip\", \"install\", \"ezkl\"])\n",
                "    subprocess.check_call([sys.executable, \"-m\", \"pip\", \"install\", \"onnx\"])\n",
                "\n",
                "# rely on local installation of ezkl if the notebook is not in colab\n",
                "except:\n",
                "    pass\n",
                "\n",
                "!RUST_LOG=trace"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "# Create TicTacToe Game States\n",
                "\n",
                "The game states are in the form of\n",
                "```json\n",
                "[\n",
                "    {\n",
                "        \"history\": [\n",
                "            [null, null, null, null, null, null, null, null, null],\n",
                "            [\"X\", null, null, null, null, null, null, null, null],\n",
                "            ...\n",
                "        ],\n",
                "        \"outcome\": \"X\"\n",
                "    }\n",
                "]\n",
                "```\n",
                "\n",
                "To create the game states for tic tac toe, do a recursive a tree search of possible gameplay."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "import json\n",
                "\n",
                "def check_winner(board):\n",
                "    winning_combinations = [\n",
                "        [0, 1, 2], [3, 4, 5], [6, 7, 8],  # rows\n",
                "        [0, 3, 6], [1, 4, 7], [2, 5, 8],  # columns\n",
                "        [0, 4, 8], [2, 4, 6]              # diagonals\n",
                "    ]\n",
                "    for combo in winning_combinations:\n",
                "        if board[combo[0]] == board[combo[1]] == board[combo[2]] and board[combo[0]] is not None:\n",
                "            return board[combo[0]]\n",
                "    return None\n",
                "\n",
                "def generate_games(board, player):\n",
                "    winner = check_winner(board)\n",
                "    if winner or None not in board:\n",
                "        # Game is over, save the outcome and board state\n",
                "        return [{\n",
                "            \"history\": [list(board)],\n",
                "            \"outcome\": winner if winner else \"Draw\"\n",
                "        }]\n",
                "\n",
                "    games = []\n",
                "    for i in range(9):\n",
                "        if board[i] is None:\n",
                "            new_board = board.copy()\n",
                "            new_board[i] = player\n",
                "            next_player = 'O' if player == 'X' else 'X'\n",
                "            next_games = generate_games(new_board, next_player)\n",
                "            for game in next_games:\n",
                "                game[\"history\"].insert(0, list(board))\n",
                "            games.extend(next_games)\n",
                "\n",
                "    return games\n",
                "\n",
                "initial_board = [None for _ in range(9)]\n",
                "games = generate_games(initial_board, 'X')\n",
                "\n",
                "with open(\"tic_tac_toe_games.json\", \"w\") as file:\n",
                "    file.write(\"[\\n\")  # Start of the list\n",
                "    for i, game in enumerate(games):\n",
                "        json.dump(game, file, separators=(',', ': '))\n",
                "        if i != len(games) - 1:  # If it's not the last game, add a comma\n",
                "            file.write(\",\\n\")\n",
                "        else:\n",
                "            file.write(\"\\n\")\n",
                "    file.write(\"]\\n\")"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "Now, let's make a list of illegal game play by running through the legal games and making bad games."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "with open(\"tic_tac_toe_games_bad.json\", \"w\") as file:\n",
                "    file.write(\"[\\n\")\n",
                "    for i, game in enumerate(games):\n",
                "        # we the permute the games\n",
                "        game_history = game['history']\n",
                "        new_game_history = []\n",
                "        for moves in game_history:\n",
                "            new_moves = []\n",
                "            for move in moves:\n",
                "                if move is None:\n",
                "                    new_moves.append('X')\n",
                "                elif move == 'X':\n",
                "                    new_moves.append('O')\n",
                "                else:\n",
                "                    new_moves.append(None)\n",
                "            new_game_history.append(new_moves)\n",
                "        game['history'] = new_game_history\n",
                "\n",
                "        json.dump(game, file, separators=(',', ': '))\n",
                "\n",
                "        if i != len(games) - 1:  # If it's not the last game, add a comma\n",
                "            file.write(\",\\n\")\n",
                "        else:\n",
                "            file.write(\"\\n\")\n",
                "    file.write(\"]\\n\")\n"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## Create a neural net to verify the execution of the tic tac toe model\n",
                "\n",
                "1. We use an autoencoder to verify the execution. An autoencoder is essentially a fan-in, fan-out architecture. It can be used for generative tasks, but in this case we use it as an anomaly detection approach.\n",
                "\n",
                "2. The autoencoder helps us extract out the latent distribution of normal tic tac toe games. If there's a weird tic tac toe game, we can then reject it if the distribution is not within some threshold we want. We can quantify this distribution via the mean absolute error between the data point and the reconstructed data."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "import torch\n",
                "import torch.nn as nn\n",
                "from torch.utils.data import DataLoader, IterableDataset, random_split\n",
                "import torch.optim as optim\n",
                "import json\n",
                "import torch.nn.functional as F\n",
                "\n",
                "class TicTacToeNet(nn.Module):\n",
                "    def __init__(self):\n",
                "        super(TicTacToeNet, self).__init__()\n",
                "        self.return_x = True\n",
                "        self.dense = nn.Linear(11 * 9 * 3, 50)\n",
                "        self.dense1 = nn.Linear(50, 20)\n",
                "        self.dense2 = nn.Linear(20, 10)\n",
                "        self.dense3 = nn.Linear(10, 20)\n",
                "        self.dense4 = nn.Linear(20, 50)\n",
                "        self.dense5 = nn.Linear(50, 11 * 9 * 3)\n",
                "\n",
                "    def forward(self, x):\n",
                "        # Store original input\n",
                "        original_x = x.clone()\n",
                "\n",
                "        # Neural network operations\n",
                "        x = F.relu(self.dense(x))\n",
                "        x = F.relu(self.dense1(x))\n",
                "        x = F.relu(self.dense2(x))\n",
                "        x = F.relu(self.dense3(x))\n",
                "        x = F.relu(self.dense4(x))\n",
                "        x = self.dense5(x)  # Not applying activation on final layer\n",
                "\n",
                "        # Calculate L1 loss\n",
                "        l1_loss = torch.mean(torch.abs(x - original_x))\n",
                "\n",
                "        if self.return_x:\n",
                "            return x, l1_loss\n",
                "        else:\n",
                "            return l1_loss"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## Load all possible tictactoe games\n",
                "\n",
                "We want to load good games and bad games for the classification later."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "class JsonDataset(IterableDataset):\n",
                "    def __init__(self, file):\n",
                "        self.file_good = file\n",
                "        self.length = self.compute_length(self.file_good)\n",
                "        self.data = self.load_data(self.file_good)\n",
                "\n",
                "    def parse_json_object(self, line):\n",
                "        try:\n",
                "            return json.loads(line)\n",
                "        except json.JSONDecodeError:\n",
                "            return None\n",
                "\n",
                "    def encode_board(self, board):\n",
                "        encoding = []\n",
                "        for cell in board:\n",
                "            if cell == 'X':\n",
                "                encoding.extend([1,0,0])\n",
                "            elif cell == 'O':\n",
                "                encoding.extend([0,1,0])\n",
                "            else:\n",
                "                encoding.extend([0,0,1])\n",
                "\n",
                "        return encoding\n",
                "\n",
                "\n",
                "    def encode_outcome(self, outcome):\n",
                "        if outcome == 'X':\n",
                "            return [1,0,0]\n",
                "        elif outcome == 'O':\n",
                "            return [0,1,0]\n",
                "        else:\n",
                "            return [0,0,1]\n",
                "\n",
                "    def compute_length(self, file_good):\n",
                "        count = 0\n",
                "        with open(file_good, 'r') as f:\n",
                "            # Skip the first line (which is \"[\")\n",
                "            next(f)\n",
                "            for line in f:\n",
                "                if line.strip() not in [\",\", \"]\"]:\n",
                "                    count += 1\n",
                "\n",
                "        return count\n",
                "\n",
                "    def __len__(self):\n",
                "        return self.length\n",
                "\n",
                "    def __getitem__(self, idx):\n",
                "        padded_history, sample_outcome = self.data[idx]\n",
                "\n",
                "        return torch.tensor(padded_history, dtype=torch.float), torch.tensor(sample_outcome, dtype=torch.float)\n",
                "\n",
                "    def process_file(self, file, is_good):\n",
                "        data = []\n",
                "        with open(file, 'r') as f:\n",
                "            next(f)\n",
                "            for line in f:\n",
                "                # Remove the trailing comma for all lines except the last one (which is \"]\")\n",
                "                if line.endswith(\",\\n\"):\n",
                "                    line = line[:-2]\n",
                "                sample = self.parse_json_object(line)\n",
                "                if sample is not None:\n",
                "                    max_length = 10  # Maximum length of a Tic Tac Toe game\n",
                "                    history = sample['history']\n",
                "\n",
                "                    if len(history) == max_length:\n",
                "                        padded_history = history\n",
                "                    else:\n",
                "                        padded_history = history + [[None] * 9 for _ in range(max_length - len(history))]\n",
                "\n",
                "                    padded_history = [self.encode_board(x) for x in padded_history]\n",
                "                    sample_outcome = self.encode_outcome(sample['outcome'])\n",
                "                    sample_outcome.extend([0,0,1] * 8)\n",
                "                    padded_history.append(sample_outcome)\n",
                "\n",
                "                    if is_good:\n",
                "                        # data.append((padded_history, sample_outcome, 0))\n",
                "                        data.append((padded_history, [1, 0]))\n",
                "                    else:\n",
                "                        # data.append((padded_history, sample_outcome, 1))\n",
                "                        data.append((padded_history, [0, 1]))\n",
                "\n",
                "        return data\n",
                "\n",
                "\n",
                "    def load_data(self, file_good):\n",
                "        data_good = self.process_file(file_good, True)\n",
                "        # data_bad = self.process_file(file_bad, False)\n",
                "        # data_good.extend(data_bad)\n",
                "\n",
                "        return data_good\n",
                "\n",
                "def collate_fn(batch):\n",
                "    histories, outcomes = zip(*batch)\n",
                "\n",
                "    # Convert nested lists to tensor\n",
                "    histories_tensor = torch.tensor(histories, dtype=torch.float32)\n",
                "    outcomes_tensor = torch.tensor(outcomes, dtype=torch.int64)\n",
                "\n",
                "    return histories_tensor, outcomes_tensor\n",
                "\n",
                "dataset = JsonDataset('tic_tac_toe_games.json')\n",
                "\n",
                "total_size = len(dataset)\n",
                "train_size = int(0.8 * total_size)\n",
                "test_size = total_size - train_size\n",
                "\n",
                "train_dataset, test_dataset = random_split(dataset, [train_size, test_size])\n",
                "\n",
                "\n",
                "train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)\n",
                "test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## Train\n",
                "\n",
                "Note that during training we want to neural network to try to recreate its inputs. To do this we compare the L1Loss or Mean Absolute Error between the prediction and the input."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
                "model = TicTacToeNet().to(device)\n",
                "criterion = nn.L1Loss()\n",
                "optimizer = optim.Adam(model.parameters(), lr=0.001)\n",
                "\n",
                "# You may want to increase this\n",
                "MAX_EPOCH = 1\n",
                "\n",
                "for epoch in range(MAX_EPOCH):\n",
                "    model.train()\n",
                "\n",
                "    for history, valid in train_loader:\n",
                "\n",
                "        history = history.to(device)\n",
                "        optimizer.zero_grad()\n",
                "        history = history.view(history.size(0), -1)\n",
                "        prediction, _ = model(history)\n",
                "        loss = criterion(prediction, history)\n",
                "        loss.backward()\n",
                "        optimizer.step()\n",
                "\n",
                "        # Gather data and report\n",
                "\n",
                "    # Validation phase\n",
                "    model.eval()\n",
                "\n",
                "\n",
                "    with torch.no_grad():\n",
                "        total_count = 0\n",
                "        total_loss = 0\n",
                "        for history, valid in test_loader:\n",
                "            history = history.to(device)\n",
                "            history = history.view(history.size(0), -1)\n",
                "\n",
                "            valid = valid.to(device)\n",
                "            prediction, _ = model(history)\n",
                "            loss = criterion(prediction, history)\n",
                "\n",
                "            total_loss += loss\n",
                "\n",
                "            total_count += 1\n",
                "\n",
                "            # get values for the accuracy\n",
                "\n",
                "\n",
                "    avg_loss = total_loss / total_count\n",
                "    print(f\"Epoch {epoch + 1}/{MAX_EPOCH} - Loss: {avg_loss:.2f}\")\n",
                "    print(history[-1])\n",
                "    print(prediction[-1])\n"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "# Get the ranges for normal data\n",
                "\n",
                "We plot the range of possible losses for normal data.\n",
                "\n"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "# set model to get loss out\n",
                "model.eval()\n",
                "model.return_x = False\n",
                "\n",
                "\n",
                "loss_list = []\n",
                "\n",
                "for history, valid in test_loader:\n",
                "    history = history.to(device)\n",
                "    history = history.view(history.size(0), -1)\n",
                "    loss = model(history)\n",
                "    loss_list.append(loss.cpu().detach().numpy())\n",
                "\n"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "# visualize\n",
                "import matplotlib.pyplot as plt\n",
                "\n",
                "plt.hist(loss_list, bins=50)\n",
                "plt.xlabel(\"Normal loss\")\n",
                "plt.ylabel(\"No of examples\")\n",
                "plt.show()"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "# Get the ranges for bad data\n",
                "\n",
                "Likewise, we plot the ranges of losses for bad data."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "# load bad data\n",
                "bad_dataset = JsonDataset('tic_tac_toe_games_bad.json')\n",
                "\n",
                "bad_total_size = len(bad_dataset)\n",
                "bad_train_size = int(0.8 * bad_total_size)\n",
                "bad_test_size = bad_total_size - bad_train_size\n",
                "\n",
                "bad_train_dataset, bad_test_dataset = random_split(bad_dataset, [bad_train_size, bad_test_size])\n",
                "\n",
                "bad_train_loader = DataLoader(bad_train_dataset, batch_size=32, shuffle=True)\n",
                "bad_test_loader = DataLoader(bad_test_dataset, batch_size=32, shuffle=False)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "bad_loss_list = []\n",
                "\n",
                "for history, valid in bad_test_loader:\n",
                "    history = history.to(device)\n",
                "    history = history.view(history.size(0), -1)\n",
                "    loss = model(history)\n",
                "    bad_loss_list.append(loss.cpu().detach().numpy())"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "import matplotlib.pyplot as plt\n",
                "\n",
                "\n",
                "plt.hist(bad_loss_list, bins=50)\n",
                "plt.xlabel(\"Normal loss\")\n",
                "plt.ylabel(\"No of examples\")\n",
                "plt.show()"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "**Note:** By visual inspection we can see that if the loss is greater than `~0.200` the data is likely anomalous (this value will probably change depending on the runs). There seems to be no overlap in loss between good gameplay and bad gameplay.\n",
                "\n",
                "We could alternatively set the threshold to 4 standard deviation away from the mean of normal values. Or you can vary this according to the error tolerance required."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "import numpy as np\n",
                "threshold = np.mean(loss_list)+ 4 * np.std(loss_list)\n",
                "print(threshold)"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## Export Onnx and Data"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "# Set the model to evaluation mode\n",
                "model = model.cpu()\n",
                "model.eval()\n",
                "\n",
                "# Obtain a sample datapoint from the DataLoader\n",
                "sample_data_iter = iter(train_loader)\n",
                "sample_history, _ = next(sample_data_iter)\n",
                "\n",
                "# Move the sample datapoint to the cpu so it is the same as the model\n",
                "x = sample_history.to('cpu')\n",
                "x = x.view(x.size(0), -1)\n",
                "\n",
                "# Export the model using ONNX\n",
                "torch.onnx.export(\n",
                "    model,                        # model being run\n",
                "    x,                            # model input (or a tuple for multiple inputs)\n",
                "    \"tictactoe_network.onnx\",     # where to save the model (can be a file or file-like object)\n",
                "    export_params=True,           # store the trained parameter weights inside the model file\n",
                "    opset_version=10,             # the ONNX version to export the model to\n",
                "    do_constant_folding=True,     # whether to execute constant folding for optimization\n",
                "    input_names=['input'],        # the model's input names\n",
                "    output_names=['output'],      # the model's output names\n",
                "    dynamic_axes={\n",
                "        'input': {0: 'batch_size'},     # variable length axes\n",
                "        'output': {0: 'batch_size'}\n",
                "    }\n",
                ")\n",
                "\n",
                "data_array = ((x[0]).detach().numpy()).reshape([-1]).tolist()\n",
                "\n",
                "data = dict(input_data = [data_array])\n",
                "\n",
                "# Serialize data into file:\n",
                "json.dump(data, open(\"data.json\", 'w'))\n"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "# Prove the Tic Tac Toe Anomaly Detection Model\n",
                "\n",
                "We import the ezkl library to setup the zk system."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "import os\n",
                "import ezkl\n",
                "\n",
                "data_path = os.path.join(\"data.json\")\n",
                "model_path = os.path.join('tictactoe_network.onnx')\n",
                "compiled_model_path = os.path.join('network.ezkl')\n",
                "pk_path = os.path.join('test.pk')\n",
                "vk_path = os.path.join('test.vk')\n",
                "settings_path = os.path.join('settings.json')\n",
                "\n",
                "witness_path = os.path.join('witness.json')\n",
                "proof_path = os.path.join('proof.json')"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "res = ezkl.gen_settings(model_path, settings_path)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "# IMPORTANT NOTE: You may want to set the scale ranges and set to \"accuracy\"\n",
                "# In it's current state, it will likely truncate values\n",
                "# For testing we will just stick to resources to reduce computational costs\n",
                "# Example:\n",
                "# ezkl.calibrate_settings(data_path, model_path, settings_path, \"accuracy\", scales = [2,9])\n",
                "cal_path = os.path.join(\"calibration.json\")\n",
                "\n",
                "data_array = ((x[0:20]).detach().numpy()).reshape([-1]).tolist()\n",
                "\n",
                "data = dict(input_data = [data_array])\n",
                "\n",
                "# Serialize data into file:\n",
                "json.dump(data, open(cal_path, 'w'))\n",
                "\n",
                "\n",
                "ezkl.calibrate_settings(cal_path, model_path, settings_path, \"resources\", scales = [4])"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "ezkl.get_srs( settings_path)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "ezkl.compile_circuit(model_path, compiled_model_path, settings_path)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "ezkl.setup(\n",
                "    compiled_model_path,\n",
                "    vk_path,\n",
                "    pk_path,\n",
                "    \n",
                ")"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "import json\n",
                "\n",
                "with open(data_path, \"r\") as f:\n",
                "    data = json.load(f)\n",
                "    print(len(data['input_data'][0]))\n",
                "\n",
                "ezkl.gen_witness(data_path, compiled_model_path, witness_path)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "ezkl.mock(witness_path, compiled_model_path)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "res = ezkl.prove(\n",
                "        witness_path,\n",
                "        compiled_model_path,\n",
                "        pk_path,\n",
                "        proof_path,\n",
                "        \n",
                "        \"single\",\n",
                "    )\n",
                "\n",
                "print(res)\n",
                "assert os.path.isfile(proof_path)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "res = ezkl.verify(\n",
                "        proof_path,\n",
                "        settings_path,\n",
                "        vk_path,\n",
                "        \n",
                "    )\n",
                "\n",
                "assert res == True\n",
                "print(\"verified\")"
            ]
        }
    ],
    "metadata": {
        "accelerator": "GPU",
        "colab": {
            "gpuType": "V100",
            "provenance": []
        },
        "kernelspec": {
            "display_name": "Python 3 (ipykernel)",
            "language": "python",
            "name": "python3"
        },
        "language_info": {
            "codemirror_mode": {
                "name": "ipython",
                "version": 3
            },
            "file_extension": ".py",
            "mimetype": "text/x-python",
            "name": "python",
            "nbconvert_exporter": "python",
            "pygments_lexer": "ipython3",
            "version": "3.9.15"
        }
    },
    "nbformat": 4,
    "nbformat_minor": 4
}
