{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Mean of ERC20 transfer amounts\n",
        "\n",
        "This notebook shows how to calculate the mean of ERC20 transfer amounts, pulling data in from a Postgres database. First we install and get the necessary libraries running. \n",
        "The first of which is [e2pg](https://github.com/indexsupply/x/tree/main/docs/e2pg), which is a library that allows us to pull data from the Ethereum blockchain into a Postgres database.\n",
        "\n",
        "Make sure you install postgres if needed https://postgresapp.com/. \n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "import os\n",
        "import getpass\n",
        "\n",
        "\n",
        "# swap out for the relevant linux/amd64, darwin/arm64, darwin/amd64, windows/amd64\n",
        "os.system(\"curl -LO https://indexsupply.net/bin/main/linux/amd64/e2pg\")\n",
        "os.system(\"chmod +x e2pg\")\n",
        "\n",
        "\n",
        "os.environ[\"PG_URL\"] = \"postgresql://\" + getpass.getuser() + \":@localhost:5432/e2pg\"\n",
        "os.environ[\"RLPS_URL\"] = \"https://1.rlps.indexsupply.net\"\n",
        "\n",
        "# print the two env variables\n",
        "os.system(\"echo $PG_URL\")\n",
        "os.system(\"echo $RLPS_URL\")\n",
        "\n",
        "os.system(\"createdb -h localhost -p 5432 e2pg\")\n",
        "# equivalent of nohup ./e2pg -reset -e $RLPS_URL -pg $PG_URL &\n",
        "e2pg_process = os.system(\"nohup ./e2pg -e $RLPS_URL -pg $PG_URL &\")\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2wIAHwqH2_mo"
      },
      "source": [
        "**Import Dependencies**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9Byiv2Nc2MsK"
      },
      "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",
        "import ezkl\n",
        "import torch\n",
        "import datetime\n",
        "import pandas as pd\n",
        "import requests\n",
        "import json\n",
        "import os\n",
        "\n",
        "# import logging\n",
        "# # # uncomment for more descriptive logging \n",
        "# FORMAT = '%(levelname)s %(name)s %(asctime)-15s %(filename)s:%(lineno)d %(message)s'\n",
        "# logging.basicConfig(format=FORMAT)\n",
        "# logging.getLogger().setLevel(logging.DEBUG)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "osjj-0Ta3E8O"
      },
      "source": [
        "**Create Computational Graph**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "x1vl9ZXF3EEW",
        "outputId": "bda21d02-fe5f-4fb2-8106-f51a8e2e67aa"
      },
      "outputs": [],
      "source": [
        "from torch import nn\n",
        "import torch\n",
        "\n",
        "\n",
        "class Model(nn.Module):\n",
        "    def __init__(self):\n",
        "        super(Model, self).__init__()\n",
        "\n",
        "    # x is a time series \n",
        "    def forward(self, x):\n",
        "        return [torch.mean(x)]\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "circuit = Model()\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "x = 0.1*torch.rand(1,*[1,5], requires_grad=True)\n",
        "\n",
        "# # print(torch.__version__)\n",
        "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
        "\n",
        "print(device)\n",
        "\n",
        "circuit.to(device)\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",
        "                      \"lol.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=11,          # 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={'input' : {0 : 'batch_size'},    # variable length axes\n",
        "                                    'output' : {0 : 'batch_size'}})\n",
        "\n",
        "# export(circuit, input_shape=[1, 20])\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "E3qCeX-X5xqd"
      },
      "source": [
        "**Set Data Source and Get Data**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "6RAMplxk5xPk",
        "outputId": "bd2158fe-0c00-44fd-e632-6a3f70cdb7c9"
      },
      "outputs": [],
      "source": [
        "# make an input.json file from the df above\n",
        "input_filename = os.path.join('input.json')\n",
        "\n",
        "pg_input_file = dict(input_data = {\n",
        "  \"host\": \"localhost\",\n",
        "  # make sure you replace this with your own username\n",
        "  \"user\": getpass.getuser(),\n",
        "  \"dbname\": \"e2pg\",\n",
        "  \"password\": \"\",\n",
        "  \"query\": \"SELECT value FROM erc20_transfers ORDER BY block_number DESC LIMIT 5\",\n",
        "  \"port\": \"5432\",\n",
        "})\n",
        "\n",
        "json_formatted_str = json.dumps(pg_input_file, indent=2)\n",
        "print(json_formatted_str)\n",
        "\n",
        "\n",
        "    # Serialize data into file:\n",
        "json.dump( pg_input_file, open(input_filename, 'w' ))\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# this corresponds to 4 batches\n",
        "calibration_filename = os.path.join('calibration.json')\n",
        "\n",
        "pg_cal_file = dict(input_data = {\n",
        "  \"host\": \"localhost\",\n",
        "  # make sure you replace this with your own username\n",
        "  \"user\": getpass.getuser(),\n",
        "  \"dbname\": \"e2pg\",\n",
        "  \"password\": \"\",\n",
        "  \"query\": \"SELECT value FROM erc20_transfers ORDER BY block_number DESC LIMIT 20\",\n",
        "  \"port\": \"5432\",\n",
        "})\n",
        "\n",
        "    # Serialize data into file:\n",
        "json.dump( pg_cal_file, open(calibration_filename, 'w' ))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eLJ7oirQ_HQR"
      },
      "source": [
        "**EZKL Workflow**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "onnx_filename = os.path.join('lol.onnx')\n",
        "compiled_filename = os.path.join('lol.compiled')\n",
        "settings_filename = os.path.join('settings.json')\n",
        "\n",
        "ezkl.gen_settings(onnx_filename, settings_filename)\n",
        "\n",
        "ezkl.calibrate_settings(\n",
        "    input_filename, onnx_filename, settings_filename, \"resources\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rNw0C9QL6W88"
      },
      "outputs": [],
      "source": [
        "# setup kzg params\n",
        "params_path = os.path.join('kzg.params')\n",
        "\n",
        "res = ezkl.get_srs(params_path, settings_filename)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "\n",
        "ezkl.compile_circuit(onnx_filename, compiled_filename, settings_filename)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "4MmE9SX66_Il",
        "outputId": "16403639-66a4-4280-ac7f-6966b75de5a3"
      },
      "outputs": [],
      "source": [
        "# generate settings\n",
        "\n",
        "\n",
        "# show the settings.json\n",
        "with open(\"settings.json\") as f:\n",
        "  data = json.load(f)\n",
        "  json_formatted_str = json.dumps(data, indent=2)\n",
        "\n",
        "  print(json_formatted_str)\n",
        "\n",
        "assert os.path.exists(\"settings.json\")\n",
        "assert os.path.exists(\"input.json\")\n",
        "assert os.path.exists(\"lol.onnx\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fULvvnK7_CMb"
      },
      "outputs": [],
      "source": [
        "pk_path = os.path.join('test.pk')\n",
        "vk_path = os.path.join('test.vk')\n",
        "params_path = os.path.join('kzg.params')\n",
        "\n",
        "\n",
        "# setup the proof\n",
        "res = ezkl.setup(\n",
        "        compiled_filename,\n",
        "        vk_path,\n",
        "        pk_path,\n",
        "        params_path,\n",
        "        settings_filename,\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_filename)\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "\n",
        "witness_path = \"witness.json\"\n",
        "\n",
        "res = ezkl.gen_witness(input_filename, compiled_filename, witness_path)\n",
        "assert os.path.isfile(witness_path)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Oog3j6Kd-Wed",
        "outputId": "5839d0c1-5b43-476e-c2f8-6707de562260"
      },
      "outputs": [],
      "source": [
        "# prove the zk circuit\n",
        "# GENERATE A PROOF\n",
        "proof_path = os.path.join('test.pf')\n",
        "\n",
        "\n",
        "proof = ezkl.prove(\n",
        "        witness_path,\n",
        "        compiled_filename,\n",
        "        pk_path,\n",
        "        proof_path,\n",
        "        params_path,\n",
        "        \"single\",\n",
        "    )\n",
        "\n",
        "\n",
        "assert os.path.isfile(proof_path)\n",
        "\n",
        "# verify\n",
        "res = ezkl.verify(\n",
        "        proof_path,\n",
        "        settings_filename,\n",
        "        vk_path,\n",
        "        params_path,\n",
        "    )\n",
        "\n",
        "assert res == True\n",
        "print(\"verified\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "W7tAa-DFAtvS"
      },
      "source": [
        "# Part 2 (Using the ZK Computational Graph Onchain!)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8Ym91kaVAIB6"
      },
      "source": [
        "**Now How Do We Do It Onchain?????**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 339
        },
        "id": "fodkNgwS70FM",
        "outputId": "827b5efd-f74f-44de-c114-861b3a86daf2"
      },
      "outputs": [],
      "source": [
        "# first we need to create evm verifier\n",
        "print(vk_path)\n",
        "print(params_path)\n",
        "print(settings_filename)\n",
        "\n",
        "\n",
        "abi_path = 'test.abi'\n",
        "sol_code_path = 'test.sol'\n",
        "\n",
        "res = ezkl.create_evm_verifier(\n",
        "        vk_path,\n",
        "        params_path,\n",
        "        settings_filename,\n",
        "        sol_code_path,\n",
        "        abi_path,\n",
        "    )\n",
        "assert res == True"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Make sure anvil is running locally first\n",
        "# run with $ anvil -p 3030\n",
        "# we use the default anvil node here\n",
        "import json\n",
        "\n",
        "address_path = os.path.join(\"address.json\")\n",
        "\n",
        "res = ezkl.deploy_evm(\n",
        "    address_path,\n",
        "    sol_code_path,\n",
        "    'http://127.0.0.1:3030'\n",
        ")\n",
        "\n",
        "assert res == True\n",
        "\n",
        "with open(address_path, 'r') as file:\n",
        "    addr = file.read().rstrip()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# read the address from addr_path\n",
        "addr = None\n",
        "with open(address_path, 'r') as f:\n",
        "    addr = f.read()\n",
        "\n",
        "res = ezkl.verify_evm(\n",
        "    addr,\n",
        "    proof_path,\n",
        "    \"http://127.0.0.1:3030\"\n",
        ")\n",
        "assert res == True"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "os.system(\"killall -9 e2pg\");"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "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.13"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
