{
    "cells": [
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "# Binary Classification Tic Tac Toe Verifier\n",
                "\n",
                "We create an ML model that verifies TicTacToe Games\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"
            ]
        },
        {
            "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. Given the data generated above classify whether the tic tac toe games are valid. This approach uses a binary classification as the tic tac toe state space is fairly small. For larger state spaces, we will want to use anomaly detection based approaches."
            ]
        },
        {
            "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.dense = nn.Linear(11 * 9 * 3, 50)  # Adjusted for 11*9 input features\n",
                "        self.dense1 = nn.Linear(50, 20)\n",
                "        self.dense2 = nn.Linear(20, 10)\n",
                "        self.dense3 = nn.Linear(10, 2)\n",
                "\n",
                "    def forward(self, x):\n",
                "        x = x.view(x.size(0), -1)  # Flatten the input\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",
                "        return x"
            ]
        },
        {
            "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_good, file_bad):\n",
                "        self.file_good = file_good\n",
                "        self.file_bad = file_bad\n",
                "        self.length = self.compute_length(self.file_good, self.file_bad)\n",
                "        self.data = self.load_data(self.file_good, self.file_bad)\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, file_bad):\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",
                "        with open(file_bad, '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",
                "        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, file_bad):\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', 'tic_tac_toe_games_bad.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",
                "In this training step we'll just use the CrossEntropyLoss and have the neural network classify good games and bad games."
            ]
        },
        {
            "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.CrossEntropyLoss()\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",
                "        valid = valid.to(device)\n",
                "        optimizer.zero_grad()\n",
                "        prediction = model(history)\n",
                "        loss = criterion(prediction, valid)\n",
                "        loss.backward()\n",
                "        optimizer.step()\n",
                "\n",
                "\n",
                "    # Validation phase\n",
                "    model.eval()\n",
                "    total_count = 0\n",
                "    correct_count = 0\n",
                "\n",
                "    with torch.no_grad():\n",
                "        for history, valid in test_loader:\n",
                "            history = history.to(device)\n",
                "            valid = valid.to(device)\n",
                "            prediction = model(history)\n",
                "\n",
                "            # Convert predictions to class labels\n",
                "            _, predicted_labels = torch.max(prediction, 1)\n",
                "\n",
                "            # Convert true labels to class labels\n",
                "            _, true_class_labels = torch.max(valid, 1)\n",
                "\n",
                "            # Count correct predictions\n",
                "            correct_predictions = (predicted_labels == true_class_labels).sum().item()\n",
                "\n",
                "            total_predictions = prediction.shape[0]\n",
                "\n",
                "            # get values for the accuracy\n",
                "            correct_count += correct_predictions\n",
                "            total_count += total_predictions\n",
                "\n",
                "\n",
                "    accuracy = 100 * correct_count / total_count\n",
                "    print(f\"Epoch {epoch + 1}/{MAX_EPOCH} - Accuracy: {accuracy:.2f}%\")\n",
                "    print(history[-1])\n",
                "    print(valid[-1])\n",
                "    print(predicted_labels[-1])\n",
                "    print(true_class_labels[-1])"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "print(prediction)\n",
                "print(valid)"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "**Note:** It's reasonable for the tic tac toe model to approach 100% classification accuracy as the state space of the game is fairly small"
            ]
        },
        {
            "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",
                "\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 Classification Model"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "!RUST_LOG=trace"
            ]
        },
        {
            "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": [
                "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": "T4",
            "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
}
