{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from zkstats.core import (\n",
    "    prover_gen_settings,\n",
    "    prover_gen_proof,\n",
    "    verifier_setup,\n",
    "    verifier_verify,\n",
    ")\n",
    "from zkstats.computation import computation_to_model\n",
    "\n",
    "cwd = os.getcwd()\n",
    "\n",
    "# FIXME: fill this in with the path to your data\n",
    "data_path = f\"{cwd}/data.json\"\n",
    "\n",
    "# Paths to the output files\n",
    "output_dir = f\"{cwd}/out\"\n",
    "os.makedirs(output_dir, exist_ok=True)\n",
    "model_onnx_path = f\"{output_dir}/model.onnx\"\n",
    "compiled_model_path = f\"{output_dir}/model.compiled\"\n",
    "\n",
    "pk_path = f\"{output_dir}/model.pk\"\n",
    "vk_path = f\"{output_dir}/model.vk\"\n",
    "proof_path = f\"{output_dir}/model.pf\"\n",
    "settings_path = f\"{output_dir}/settings.json\"\n",
    "witness_path = f\"{output_dir}/witness.json\"\n",
    "comb_data_path = f\"{output_dir}/comb_data.json\""
   ]
  },
  {
   "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, x: list[torch.Tensor]):\n",
    "    out_0 = state.median(x[0])\n",
    "    out_1 = state.median(x[1])\n",
    "    return state.mean(torch.tensor([out_0, out_1]).reshape(1,-1,1))\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": 40,
   "metadata": {},
   "outputs": [],
   "source": [
    "# This is just a dummy computation. Replace it with user's computation\n",
    "#\n",
    "# import torch\n",
    "# from zkstats.computation import State\n",
    "#\n",
    "# def computation(state: State, x: list[torch.Tensor]):\n",
    "#     x_0 = x[0]\n",
    "#     out_0 = state.median(x_0)\n",
    "#     out_1 = state.median(x_0)\n",
    "#     return state.mean(torch.tensor([out_0, out_1]).reshape(1,-1,1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Generate settings and setup with user's computation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/var/folders/t3/5psrvr1x0w1_6n9kx2n7d9700000gn/T/ipykernel_97835/3906390486.py:5: TracerWarning: torch.tensor results are registered as constants in the trace. You can safely ignore this warning if you use this function to create tensors out of constant variables that would be the same every time you call this function. In any other case, this might cause the trace to be incorrect.\n",
      "  return state.mean(torch.tensor([out_0, out_1]).reshape(1,-1,1))\n",
      "/var/folders/t3/5psrvr1x0w1_6n9kx2n7d9700000gn/T/ipykernel_97835/3906390486.py:5: TracerWarning: Converting a tensor to a Python float might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!\n",
      "  return state.mean(torch.tensor([out_0, out_1]).reshape(1,-1,1))\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "==== Generate & Calibrate Setting ====\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "spawning module 0\n",
      "spawning module 2\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "scale:  default\n",
      "setting:  {\"run_args\":{\"tolerance\":{\"val\":0.0,\"scale\":1.0},\"input_scale\":8,\"param_scale\":8,\"scale_rebase_multiplier\":10,\"lookup_range\":[-25112,24986],\"logrows\":16,\"num_inner_cols\":2,\"variables\":[[\"batch_size\",1]],\"input_visibility\":{\"Hashed\":{\"hash_is_public\":true,\"outlets\":[]}},\"output_visibility\":\"Public\",\"param_visibility\":\"Private\"},\"num_rows\":14432,\"total_assignments\":12049,\"total_const_size\":1816,\"model_instance_shapes\":[[1],[1]],\"model_output_scales\":[0,8],\"model_input_scales\":[8],\"module_sizes\":{\"kzg\":[],\"poseidon\":[14432,[1]],\"elgamal\":[0,[0]]},\"required_lookups\":[\"Abs\",{\"GreaterThan\":{\"a\":0.0}},\"KroneckerDelta\"],\"check_mode\":\"UNSAFE\",\"version\":\"7.0.0\",\"num_blinding_factors\":null}\n",
      "==== setting up ezkl ====\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "spawning module 0\n",
      "spawning module 2\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Time setup: 12.489933013916016 seconds\n"
     ]
    }
   ],
   "source": [
    "_, model = computation_to_model(computation)\n",
    "prover_gen_settings(\n",
    "    [data_path],\n",
    "    comb_data_path,\n",
    "    model,\n",
    "    model_onnx_path,\n",
    "    \"default\",\n",
    "    \"resources\",\n",
    "    settings_path,\n",
    ")\n",
    "verifier_setup(model_onnx_path, compiled_model_path, settings_path, vk_path, pk_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Generate proof with your data and user's computation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "==== Generating Witness ====\n",
      "witness boolean:  1.0\n",
      "witness result 1 : 49.55078125\n",
      "==== Generating Proof ====\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "spawning module 0\n",
      "spawning module 2\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "proof:  {'instances': [[[3042937791208075219, 8157070662846698822, 3804781648660056856, 172406108020799675], [12436184717236109307, 3962172157175319849, 7381016538464732718, 1011752739694698287], [18341455175509539295, 12796101019039945164, 1607286914885633240, 1929881192315725821]]], 'proof': '[encoded data, 8960 characters]', 'transcript_type': 'EVM'}\n",
      "Time gen prf: 17.092010021209717 seconds\n"
     ]
    }
   ],
   "source": [
    "prover_gen_proof(\n",
    "    model_onnx_path,\n",
    "    comb_data_path,\n",
    "    witness_path,\n",
    "    compiled_model_path,\n",
    "    settings_path,\n",
    "    proof_path,\n",
    "    pk_path,\n",
    ")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Verify the proof to ensure it is correct"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "!@# res:  True\n",
      "prf instances:  [[[3042937791208075219, 8157070662846698822, 3804781648660056856, 172406108020799675], [12436184717236109307, 3962172157175319849, 7381016538464732718, 1011752739694698287], [18341455175509539295, 12796101019039945164, 1607286914885633240, 1929881192315725821]]]\n",
      "num_inputs:  1\n",
      "proof boolean:  1.0\n",
      "proof result 0 : 49.55078125\n"
     ]
    }
   ],
   "source": [
    "verifier_verify(proof_path, settings_path, vk_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": 44,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model onnx:\t\t /Users/mhchia/projects/work/pse/demo-next/public/assets/out/model.onnx\n",
      "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/model.vk\n"
     ]
    }
   ],
   "source": [
    "model_onnx_path, settings_path, proof_path, vk_path\n",
    "print(\"Model onnx:\\t\\t\", model_onnx_path)\n",
    "print(\"Settings:\\t\\t\", settings_path)\n",
    "print(\"Proof:\\t\\t\\t\", proof_path)\n",
    "print(\"Verification key:\\t\", vk_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.9.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
