{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "cf69bb3f-94e6-4dba-92cd-ce08df117d67",
   "metadata": {},
   "source": [
    "## World rotation\n",
    "\n",
    "Here we demonstrate how to use the EZKL package to rotate an on-chain world. \n",
    "\n",
    "![zk-gaming-diagram-transformed](https://hackmd.io/_uploads/HkApuQGV6.png)\n",
    "> **A typical ZK application flow**. For the shape rotators out there — this is an easily digestible example. A user computes a ZK-proof that they have calculated a valid rotation of a world. They submit this proof to a verifier contract which governs an on-chain world, along with a new set of coordinates, and the world rotation updates. Observe that it’s possible for one player to initiate a *global* change.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "95613ee9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# check if notebook is in colab\n",
    "try:\n",
    "    # install ezkl\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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch import nn\n",
    "import ezkl\n",
    "import os\n",
    "import json\n",
    "import torch\n",
    "import math\n",
    "\n",
    "# these are constants for the rotation\n",
    "phi = torch.tensor(5 * math.pi / 180)\n",
    "s = torch.sin(phi)\n",
    "c = torch.cos(phi)\n",
    "\n",
    "\n",
    "class RotateStuff(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(RotateStuff, self).__init__()\n",
    "\n",
    "        # create a rotation matrix -- the matrix is constant and is transposed for convenience\n",
    "        self.rot = torch.stack([torch.stack([c, -s]),\n",
    "                   torch.stack([s, c])]).t()\n",
    "\n",
    "    def forward(self, x):\n",
    "        x_rot = x @ self.rot   # same as x_rot = (rot @ x.t()).t() due to rot in O(n) (SO(n) even)\n",
    "        return x_rot\n",
    "\n",
    "\n",
    "circuit = RotateStuff()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This will showcase the principle directions of rotation by plotting the rotation of a single unit vector."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib import pyplot\n",
    "pyplot.figure(figsize=(3, 3))\n",
    "pyplot.arrow(0, 0, 1, 0, width=0.02, alpha=0.5)\n",
    "pyplot.arrow(0, 0, 0, 1, width=0.02, alpha=0.5)\n",
    "pyplot.arrow(0, 0, circuit.rot[0, 0].item(), circuit.rot[0, 1].item(), width=0.02)\n",
    "pyplot.arrow(0, 0, circuit.rot[1, 0].item(), circuit.rot[1, 1].item(), width=0.02)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b37637c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_path = os.path.join('network.onnx')\n",
    "compiled_model_path = os.path.join('network.compiled')\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",
    "srs_path = os.path.join('kzg.srs')\n",
    "witness_path = os.path.join('witness.json')\n",
    "data_path = os.path.join('input.json')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82db373a",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "# initial principle vectors for the rotation are as in the plot above\n",
    "x = torch.tensor([[1, 0], [0, 1]], dtype=torch.float32)\n",
    "\n",
    "# Flips the neural net into inference mode\n",
    "circuit.eval()\n",
    "\n",
    "    # Export the model\n",
    "torch.onnx.export(circuit,               # model being run\n",
    "                      x,                   # model input (or a tuple for multiple inputs)\n",
    "                      model_path,            # 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",
    "                      )\n",
    "\n",
    "data_array = ((x).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_path, 'w' ))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### World rotation in 2D on-chain"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For demo purposes we deploy these coordinates to a contract running locally using Anvil. This creates our on-chain world. We then rotate the world using the EZKL package and submit the proof to the contract. The contract then updates the world rotation. For demo purposes we do this repeatedly, rotating the world by 1 transform each time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import subprocess\n",
    "import time\n",
    "import threading\n",
    "\n",
    "# make sure anvil is running locally\n",
    "# $ anvil -p 3030\n",
    "\n",
    "RPC_URL = \"http://localhost:3030\"\n",
    "\n",
    "# Save process globally\n",
    "anvil_process = None\n",
    "\n",
    "def start_anvil():\n",
    "    global anvil_process\n",
    "    if anvil_process is None:\n",
    "        anvil_process = subprocess.Popen([\"anvil\", \"-p\", \"3030\", \"--code-size-limit=41943040\"])\n",
    "        if anvil_process.returncode is not None:\n",
    "            raise Exception(\"failed to start anvil process\")\n",
    "        time.sleep(3)\n",
    "\n",
    "def stop_anvil():\n",
    "    global anvil_process\n",
    "    if anvil_process is not None:\n",
    "        anvil_process.terminate()\n",
    "        anvil_process = None\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We define our `PyRunArgs` objects which contains the visibility parameters for out model. \n",
    "- `input_visibility` defines the visibility of the model inputs\n",
    "- `param_visibility` defines the visibility of the model weights and constants and parameters \n",
    "- `output_visibility` defines the visibility of the model outputs\n",
    "\n",
    "Here we create the following setup:\n",
    "- `input_visibility`: \"public\"\n",
    "- `param_visibility`: \"fixed\"\n",
    "- `output_visibility`: public"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d5e374a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "py_run_args = ezkl.PyRunArgs()\n",
    "py_run_args.input_visibility = \"public\"\n",
    "py_run_args.output_visibility = \"public\"\n",
    "py_run_args.param_visibility = \"private\" # private by default\n",
    "py_run_args.scale_rebase_multiplier = 10\n",
    "\n",
    "res = ezkl.gen_settings(model_path, settings_path, py_run_args=py_run_args)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3aa4f090",
   "metadata": {},
   "outputs": [],
   "source": [
    "res = ezkl.compile_circuit(model_path, compiled_model_path, settings_path)\n",
    "assert res == True"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We also define a contract that holds out test data. This contract will contain in its storage the data that we will read from and attest to. In production you would not need to set up a local anvil instance. Instead you would replace RPC_URL with the actual RPC endpoint of the chain you are deploying your verifiers too, reading from the data on said chain."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ezkl.setup_test_evm_witness(\n",
    "    data_path,\n",
    "    compiled_model_path,\n",
    "    # we write the call data to the same file as the input data\n",
    "    data_path,\n",
    "    input_source=ezkl.PyTestDataSource.OnChain,\n",
    "    output_source=ezkl.PyTestDataSource.File,\n",
    "    rpc_url=RPC_URL)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As we use Halo2 with KZG-commitments we need an SRS string from (preferably) a multi-party trusted setup ceremony. For an overview of the procedures for such a ceremony check out [this page](https://blog.ethereum.org/2023/01/16/announcing-kzg-ceremony). The `get_srs` command retrieves a correctly sized SRS given the calibrated settings file from [here](https://github.com/han0110/halo2-kzg-srs). \n",
    "\n",
    "These SRS were generated with [this](https://github.com/privacy-scaling-explorations/perpetualpowersoftau) ceremony. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8b74dcee",
   "metadata": {},
   "outputs": [],
   "source": [
    "# srs path\n",
    "res = ezkl.get_srs( settings_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18c8b7c7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# now generate the witness file \n",
    "\n",
    "witness = ezkl.gen_witness(data_path, compiled_model_path, witness_path)\n",
    "assert os.path.isfile(witness_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here we setup verifying and proving keys for the circuit. As the name suggests the proving key is needed for ... proving and the verifying key is needed for ... verifying. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b1c561a8",
   "metadata": {},
   "outputs": [],
   "source": [
    "res = ezkl.setup(\n",
    "        compiled_model_path,\n",
    "        vk_path,\n",
    "        pk_path,\n",
    "        \n",
    "    )\n",
    "\n",
    "assert res == True\n",
    "assert os.path.isfile(vk_path)\n",
    "assert os.path.isfile(pk_path)\n",
    "assert os.path.isfile(settings_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now create an EVM verifier contract from our circuit. This contract will be deployed to the chain we are using. In this case we are using a local anvil instance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "abi_path = 'test.abi'\n",
    "sol_code_path = 'test.sol'\n",
    "\n",
    "res = ezkl.create_evm_verifier(\n",
    "        vk_path,\n",
    "        \n",
    "        settings_path,\n",
    "        sol_code_path,\n",
    "        abi_path,\n",
    "    )\n",
    "assert res == True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "addr_path_verifier = \"addr_verifier.txt\"\n",
    "\n",
    "res = ezkl.deploy_evm(\n",
    "    addr_path_verifier,\n",
    "    sol_code_path,\n",
    "    'http://127.0.0.1:3030'\n",
    ")\n",
    "\n",
    "assert res == True"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With the vanilla verifier deployed, we can now create the data attestation contract, which will read in the instances from the calldata to the verifier, attest to them, call the verifier and then return the result. \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "abi_path = 'test.abi'\n",
    "sol_code_path = 'test.sol'\n",
    "input_path = 'input.json'\n",
    "\n",
    "res = ezkl.create_evm_data_attestation(\n",
    "        input_path,\n",
    "        settings_path,\n",
    "        sol_code_path,\n",
    "        abi_path,\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "addr_path_da = \"addr_da.txt\"\n",
    "\n",
    "res = ezkl.deploy_da_evm(\n",
    "        addr_path_da,\n",
    "        input_path,\n",
    "        settings_path,\n",
    "        sol_code_path,\n",
    "        RPC_URL,\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can pull in the data from the contract and calculate a new set of coordinates. We then rotate the world by 1 transform and submit the proof to the contract. The contract could then update the world rotation (logic not inserted here). For demo purposes we do this repeatedly, rotating the world by 1 transform. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c384cbc8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# GENERATE A PROOF\n",
    "\n",
    "\n",
    "proof_path = os.path.join('test.pf')\n",
    "\n",
    "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": "markdown",
   "metadata": {},
   "source": [
    "Call the view only verify method on the contract to verify the proof. Since it is a view function this is safe to use in production since you don't have to pass your private key."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76f00d41",
   "metadata": {},
   "outputs": [],
   "source": [
    "# read the verifier address\n",
    "addr_verifier = None\n",
    "with open(addr_path_verifier, 'r') as f:\n",
    "    addr = f.read()\n",
    "#read the data attestation address\n",
    "addr_da = None\n",
    "with open(addr_path_da, 'r') as f:\n",
    "    addr_da = f.read()\n",
    "\n",
    "res = ezkl.verify_evm(\n",
    "    addr,\n",
    "    proof_path,\n",
    "    RPC_URL,\n",
    "    addr_da,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As a sanity check lets plot the rotations of the unit vectors. We can see that the unit vectors rotate as expected by the output of the circuit. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "witness['outputs'][0][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "settings = json.load(open(settings_path, 'r'))\n",
    "out_scale = settings[\"model_output_scales\"][0]\n",
    "\n",
    "from matplotlib import pyplot\n",
    "pyplot.figure(figsize=(3, 3))\n",
    "pyplot.arrow(0, 0, 1, 0, width=0.02, alpha=0.5)\n",
    "pyplot.arrow(0, 0, 0, 1, width=0.02, alpha=0.5)\n",
    "\n",
    "arrow_x = ezkl.string_to_float(witness['outputs'][0][0], out_scale)\n",
    "arrow_y = ezkl.string_to_float(witness['outputs'][0][1], out_scale)\n",
    "pyplot.arrow(0, 0, arrow_x, arrow_y, width=0.02)\n",
    "arrow_x = ezkl.string_to_float(witness['outputs'][0][2], out_scale)\n",
    "arrow_y = ezkl.string_to_float(witness['outputs'][0][3], out_scale)\n",
    "pyplot.arrow(0, 0, arrow_x, arrow_y, width=0.02)"
   ]
  }
 ],
 "metadata": {
  "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": 5
}
