{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: torch in /Users/mhchia/Library/Caches/pypoetry/virtualenvs/zkstats-brXmXluj-py3.12/lib/python3.12/site-packages (2.2.0)\n",
      "Requirement already satisfied: zkstats==0.1.10 in /Users/mhchia/Library/Caches/pypoetry/virtualenvs/zkstats-brXmXluj-py3.12/lib/python3.12/site-packages (0.1.10)\n",
      "Requirement already satisfied: click<9.0.0,>=8.1.7 in /Users/mhchia/Library/Caches/pypoetry/virtualenvs/zkstats-brXmXluj-py3.12/lib/python3.12/site-packages (from zkstats==0.1.10) (8.1.7)\n",
      "Requirement already satisfied: ezkl==9.1.0 in /Users/mhchia/Library/Caches/pypoetry/virtualenvs/zkstats-brXmXluj-py3.12/lib/python3.12/site-packages (from zkstats==0.1.10) (9.1.0)\n",
      "Requirement already satisfied: matplotlib<4.0.0,>=3.8.2 in /Users/mhchia/Library/Caches/pypoetry/virtualenvs/zkstats-brXmXluj-py3.12/lib/python3.12/site-packages (from zkstats==0.1.10) (3.9.2)\n",
      "Requirement already satisfied: numpy<2.0.0,>=1.26.2 in /Users/mhchia/Library/Caches/pypoetry/virtualenvs/zkstats-brXmXluj-py3.12/lib/python3.12/site-packages (from zkstats==0.1.10) (1.26.4)\n",
      "Requirement already satisfied: onnx<2.0.0,>=1.15.0 in /Users/mhchia/Library/Caches/pypoetry/virtualenvs/zkstats-brXmXluj-py3.12/lib/python3.12/site-packages (from zkstats==0.1.10) (1.16.2)\n",
      "Requirement already satisfied: requests<3.0.0,>=2.31.0 in /Users/mhchia/Library/Caches/pypoetry/virtualenvs/zkstats-brXmXluj-py3.12/lib/python3.12/site-packages (from zkstats==0.1.10) (2.32.3)\n",
      "Requirement already satisfied: scipy<2.0.0,>=1.11.4 in /Users/mhchia/Library/Caches/pypoetry/virtualenvs/zkstats-brXmXluj-py3.12/lib/python3.12/site-packages (from zkstats==0.1.10) (1.13.1)\n",
      "Requirement already satisfied: statistics<2.0.0,>=1.0.3 in /Users/mhchia/Library/Caches/pypoetry/virtualenvs/zkstats-brXmXluj-py3.12/lib/python3.12/site-packages (from zkstats==0.1.10) (1.0.3.5)\n",
      "Requirement already satisfied: filelock in /Users/mhchia/Library/Caches/pypoetry/virtualenvs/zkstats-brXmXluj-py3.12/lib/python3.12/site-packages (from torch) (3.16.0)\n",
      "Requirement already satisfied: typing-extensions>=4.8.0 in /Users/mhchia/Library/Caches/pypoetry/virtualenvs/zkstats-brXmXluj-py3.12/lib/python3.12/site-packages (from torch) (4.12.2)\n",
      "Requirement already satisfied: sympy in /Users/mhchia/Library/Caches/pypoetry/virtualenvs/zkstats-brXmXluj-py3.12/lib/python3.12/site-packages (from torch) (1.13.2)\n",
      "Requirement already satisfied: networkx in /Users/mhchia/Library/Caches/pypoetry/virtualenvs/zkstats-brXmXluj-py3.12/lib/python3.12/site-packages (from torch) (3.2.1)\n",
      "Requirement already satisfied: jinja2 in /Users/mhchia/Library/Caches/pypoetry/virtualenvs/zkstats-brXmXluj-py3.12/lib/python3.12/site-packages (from torch) (3.1.4)\n",
      "Requirement already satisfied: fsspec in /Users/mhchia/Library/Caches/pypoetry/virtualenvs/zkstats-brXmXluj-py3.12/lib/python3.12/site-packages (from torch) (2024.9.0)\n",
      "Requirement already satisfied: contourpy>=1.0.1 in /Users/mhchia/Library/Caches/pypoetry/virtualenvs/zkstats-brXmXluj-py3.12/lib/python3.12/site-packages (from matplotlib<4.0.0,>=3.8.2->zkstats==0.1.10) (1.3.0)\n",
      "Requirement already satisfied: cycler>=0.10 in /Users/mhchia/Library/Caches/pypoetry/virtualenvs/zkstats-brXmXluj-py3.12/lib/python3.12/site-packages (from matplotlib<4.0.0,>=3.8.2->zkstats==0.1.10) (0.12.1)\n",
      "Requirement already satisfied: fonttools>=4.22.0 in /Users/mhchia/Library/Caches/pypoetry/virtualenvs/zkstats-brXmXluj-py3.12/lib/python3.12/site-packages (from matplotlib<4.0.0,>=3.8.2->zkstats==0.1.10) (4.53.1)\n",
      "Requirement already satisfied: kiwisolver>=1.3.1 in /Users/mhchia/Library/Caches/pypoetry/virtualenvs/zkstats-brXmXluj-py3.12/lib/python3.12/site-packages (from matplotlib<4.0.0,>=3.8.2->zkstats==0.1.10) (1.4.7)\n",
      "Requirement already satisfied: packaging>=20.0 in /Users/mhchia/Library/Caches/pypoetry/virtualenvs/zkstats-brXmXluj-py3.12/lib/python3.12/site-packages (from matplotlib<4.0.0,>=3.8.2->zkstats==0.1.10) (24.1)\n",
      "Requirement already satisfied: pillow>=8 in /Users/mhchia/Library/Caches/pypoetry/virtualenvs/zkstats-brXmXluj-py3.12/lib/python3.12/site-packages (from matplotlib<4.0.0,>=3.8.2->zkstats==0.1.10) (10.4.0)\n",
      "Requirement already satisfied: pyparsing>=2.3.1 in /Users/mhchia/Library/Caches/pypoetry/virtualenvs/zkstats-brXmXluj-py3.12/lib/python3.12/site-packages (from matplotlib<4.0.0,>=3.8.2->zkstats==0.1.10) (3.1.4)\n",
      "Requirement already satisfied: python-dateutil>=2.7 in /Users/mhchia/Library/Caches/pypoetry/virtualenvs/zkstats-brXmXluj-py3.12/lib/python3.12/site-packages (from matplotlib<4.0.0,>=3.8.2->zkstats==0.1.10) (2.9.0.post0)\n",
      "Requirement already satisfied: protobuf>=3.20.2 in /Users/mhchia/Library/Caches/pypoetry/virtualenvs/zkstats-brXmXluj-py3.12/lib/python3.12/site-packages (from onnx<2.0.0,>=1.15.0->zkstats==0.1.10) (5.28.1)\n",
      "Requirement already satisfied: charset-normalizer<4,>=2 in /Users/mhchia/Library/Caches/pypoetry/virtualenvs/zkstats-brXmXluj-py3.12/lib/python3.12/site-packages (from requests<3.0.0,>=2.31.0->zkstats==0.1.10) (3.3.2)\n",
      "Requirement already satisfied: idna<4,>=2.5 in /Users/mhchia/Library/Caches/pypoetry/virtualenvs/zkstats-brXmXluj-py3.12/lib/python3.12/site-packages (from requests<3.0.0,>=2.31.0->zkstats==0.1.10) (3.8)\n",
      "Requirement already satisfied: urllib3<3,>=1.21.1 in /Users/mhchia/Library/Caches/pypoetry/virtualenvs/zkstats-brXmXluj-py3.12/lib/python3.12/site-packages (from requests<3.0.0,>=2.31.0->zkstats==0.1.10) (2.2.3)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /Users/mhchia/Library/Caches/pypoetry/virtualenvs/zkstats-brXmXluj-py3.12/lib/python3.12/site-packages (from requests<3.0.0,>=2.31.0->zkstats==0.1.10) (2024.8.30)\n",
      "Requirement already satisfied: docutils>=0.3 in /Users/mhchia/Library/Caches/pypoetry/virtualenvs/zkstats-brXmXluj-py3.12/lib/python3.12/site-packages (from statistics<2.0.0,>=1.0.3->zkstats==0.1.10) (0.21.2)\n",
      "Requirement already satisfied: MarkupSafe>=2.0 in /Users/mhchia/Library/Caches/pypoetry/virtualenvs/zkstats-brXmXluj-py3.12/lib/python3.12/site-packages (from jinja2->torch) (2.1.5)\n",
      "Requirement already satisfied: mpmath<1.4,>=1.1.0 in /Users/mhchia/Library/Caches/pypoetry/virtualenvs/zkstats-brXmXluj-py3.12/lib/python3.12/site-packages (from sympy->torch) (1.3.0)\n",
      "Requirement already satisfied: six>=1.5 in /Users/mhchia/Library/Caches/pypoetry/virtualenvs/zkstats-brXmXluj-py3.12/lib/python3.12/site-packages (from python-dateutil>=2.7->matplotlib<4.0.0,>=3.8.2->zkstats==0.1.10) (1.16.0)\n",
      "Note: you may need to restart the kernel to use updated packages.\n"
     ]
    }
   ],
   "source": [
    "%pip install torch zkstats==0.1.10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from zkstats.core import (\n",
    "    setup,\n",
    ")\n",
    "from zkstats.computation import computation_to_model\n",
    "\n",
    "cwd = os.getcwd()\n",
    "\n",
    "\n",
    "# Paths to the output files\n",
    "output_dir = f\"{cwd}/out\"\n",
    "os.makedirs(output_dir, exist_ok=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Shared by the data provider beforehand\n",
    "data_shape = {'x': 7, 'y': 7}\n",
    "data_commitment_path = f\"{output_dir}/data_commitment.json\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "User select the columns"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## User-defined Computation\n",
    "\n",
    "A computation should be of type `TComputation`. For example, the following code snippet defines a computation that computes the sum of the private data.\n",
    "\n",
    "```python\n",
    "def computation(state: State, args: Args):\n",
    "    x = args[\"x\"]\n",
    "    y = args[\"y\"]\n",
    "    return state.mean(x), state.mean(y)\n",
    "```\n",
    "\n",
    "FIXME: The following code snippet is entirely from the user. You MUST check\n",
    "1. the code only performs zkstats-related operations.\n",
    "2. the computation must not leak any information about the private data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "# This is just a dummy computation. Replace it with user's computation\n",
    "import torch\n",
    "from zkstats.computation import State, Args\n",
    "\n",
    "def computation(state: State, args: Args):\n",
    "    x = args[\"x\"]\n",
    "    y = args[\"y\"]\n",
    "    return state.mean(x), state.mean(y)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Get proof, settings, and precal_witness from the data provider"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "proof_path = f\"{output_dir}/model.pf\"\n",
    "settings_path = f\"{output_dir}/settings.json\"\n",
    "precal_witness_path = f\"{output_dir}/precal_witness.json\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Verify the proof to ensure it is correct\n",
    "NOTE: The following section is to illustrate what should be done on the user (data consumer) side. This step is not required by the data provider."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "==== setting up ezkl ====\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/mhchia/Library/Caches/pypoetry/virtualenvs/zkstats-brXmXluj-py3.12/lib/python3.12/site-packages/torch/onnx/symbolic_opset9.py:2174: FutureWarning: 'torch.onnx.symbolic_opset9._cast_Bool' is deprecated in version 2.0 and will be removed in the future. Please Avoid using this function and create a Cast node instead.\n",
      "  return fn(g, to_cast_func(g, input, False), to_cast_func(g, other, False))\n",
      "/Users/mhchia/Library/Caches/pypoetry/virtualenvs/zkstats-brXmXluj-py3.12/lib/python3.12/site-packages/torch/onnx/utils.py:1703: UserWarning: The exported ONNX model failed ONNX shape inference. The model will not be executable by the ONNX Runtime. If this is unintended and you believe there is a bug, please report an issue at https://github.com/pytorch/pytorch/issues. Error reported by strict ONNX shape inference: [ShapeInferenceError] (op_type:Where, node name: /Where_1): Y has inconsistent type tensor(float) (Triggered internally at /Users/runner/work/pytorch/pytorch/pytorch/torch/csrc/jit/serialization/export.cpp:1490.)\n",
      "  _C._check_onnx_proto(proto)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Time setup: 0.907416820526123 seconds\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[51.5390625, 4.0859375]"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import json\n",
    "import numpy as np\n",
    "from typing import Dict\n",
    "\n",
    "\n",
    "def create_dummy(shape_info: Dict[str, int], dummy_data_path: str) -> None:\n",
    "    \"\"\"\n",
    "    Create a dummy data file with randomized data based on the provided shape information.\n",
    "\n",
    "    Parameters:\n",
    "    - shape_info (dict): A dictionary where keys are column names and values are the number of elements (shape).\n",
    "    - dummy_data_path (str): The path to save the dummy data file.\n",
    "    \"\"\"\n",
    "    dummy_data = {}\n",
    "    for col, length in shape_info.items():\n",
    "        # Generate random data for each column\n",
    "        dummy_data[col] = np.round(np.random.uniform(0, 100, length), 1).tolist()\n",
    "\n",
    "    with open(dummy_data_path, 'w') as f:\n",
    "        json.dump(dummy_data, f)\n",
    "\n",
    "\n",
    "from zkstats.core import verifier_define_calculation, verifier_verify\n",
    "\n",
    "\n",
    "def generate_verifier_model_and_setup(\n",
    "    output_dir: str,\n",
    "    computation,  # from myself\n",
    "    precal_witness_path,  # from prover\n",
    "    data_shape,  # from prover\n",
    "):\n",
    "    # Generate the verifier model with the `precal_witness_path` provided by the prover\n",
    "    selected_columns, _, verifier_model = computation_to_model(computation, precal_witness_path, data_shape, isProver=False)\n",
    "\n",
    "    verifier_model_path = f\"{output_dir}/verifier_model.onnx\"\n",
    "    verifier_compiled_model_path = f\"{output_dir}/verifier_model.compiled\"\n",
    "\n",
    "    verifier_vk_path = f\"{output_dir}/verifier_model.vk\"\n",
    "    verifier_pk_path = f\"{output_dir}/verifier_model.pk\"\n",
    "    dummy_data_path = f\"{output_dir}/dummy_data.json\"\n",
    "    sel_dummy_data_path = f\"{output_dir}/sel_dummy_data.json\"\n",
    "\n",
    "    # Create dummy data with the same shape as the original data\n",
    "    create_dummy(data_shape, dummy_data_path)\n",
    "\n",
    "    # Generate the verifier model given the dummy data and the selected columns\n",
    "    verifier_define_calculation(dummy_data_path, selected_columns, sel_dummy_data_path, verifier_model, verifier_model_path)\n",
    "\n",
    "    # Generate the verification key\n",
    "    setup(verifier_model_path, verifier_compiled_model_path, settings_path, verifier_vk_path, verifier_pk_path)\n",
    "\n",
    "    return selected_columns, verifier_vk_path\n",
    "\n",
    "# Determine which srs to use with the logrows in the settings.json\n",
    "with open(settings_path, \"r\") as f:\n",
    "    settings = json.load(f)\n",
    "logrows = int(settings[\"run_args\"][\"logrows\"])\n",
    "srs_path = f'~/.ezkl/srs/kzg{logrows}.srs'\n",
    "\n",
    "# Call the function with the necessary parameters\n",
    "selected_columns, verifier_vk_path = generate_verifier_model_and_setup(\n",
    "    output_dir,\n",
    "    computation,\n",
    "    precal_witness_path,\n",
    "    data_shape,\n",
    ")\n",
    "\n",
    "# verify the proof\n",
    "verifier_verify(proof_path, settings_path, verifier_vk_path, selected_columns, data_commitment_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Print the file paths. You should share the following files back to the user for them to verify the proof. You **SHOULD NOT** share more files otherwise data might be leaked."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Settings:\t\t /Users/mhchia/projects/work/pse/demo-next/public/assets/out/settings.json\n",
      "Proof:\t\t\t /Users/mhchia/projects/work/pse/demo-next/public/assets/out/model.pf\n",
      "Verification key:\t /Users/mhchia/projects/work/pse/demo-next/public/assets/out/verifier_model.vk\n",
      "Srs path:\t\t ~/.ezkl/srs/kzg12.srs\n"
     ]
    }
   ],
   "source": [
    "print(\"Settings:\\t\\t\", settings_path)\n",
    "print(\"Proof:\\t\\t\\t\", proof_path)\n",
    "print(\"Verification key:\\t\", verifier_vk_path)\n",
    "print(\"Srs path:\\t\\t\", srs_path)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.12.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
