{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "805d746e-610b-4d40-80d2-a8080a993f96",
   "metadata": {},
   "source": [
    "# Simulating EPA-RE using points of low-order"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b4386513-cc14-434b-a748-2863f8657452",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "import itertools\n",
    "import glob\n",
    "\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "from collections import Counter\n",
    "\n",
    "from pathlib import Path\n",
    "from random import randint, randbytes\n",
    "from typing import Type, Any\n",
    "\n",
    "from bs4 import BeautifulSoup\n",
    "from tqdm.auto import tqdm, trange\n",
    "\n",
    "from pyecsca.ec.params import DomainParameters, get_params\n",
    "from pyecsca.ec.mult import *\n",
    "from pyecsca.sca.re.rpa import multiples_computed\n",
    "from pyecsca.misc.utils import TaskExecutor\n",
    "\n",
    "from common import *"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5b156d2a-7345-47f8-a76e-71a7d2be9d22",
   "metadata": {},
   "source": [
    "## Initialize"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a660e3ac-401b-47a0-92de-55afe63c420a",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(len(all_mults))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a95b27fc-96a9-41b5-9972-dc8386ed386d",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(len(all_mults_with_ctr))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "07bc266d-35eb-4f6d-bdba-e9f6f66827f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Needs imports on the inside to be spawn enabled to save memory.\n",
    "\n",
    "def get_general_multiples(bits: int, samples: int = 1000) -> MultResults:\n",
    "    from random import randint\n",
    "    results = []\n",
    "    for _ in range(samples):\n",
    "        big_scalar = randint(1, 2**bits)\n",
    "        results.append({big_scalar})\n",
    "    return MultResults(results, samples)\n",
    "\n",
    "def get_general_n_multiples(bits: int, n: int, samples: int = 1000) -> MultResults:\n",
    "    from random import randint\n",
    "    results = []\n",
    "    for _ in range(samples):\n",
    "        smult = set()\n",
    "        for i in range(n):\n",
    "            b = randint(1,256)\n",
    "            smult.add(randint(2**b,2**(b+1)))\n",
    "        results.append(smult)\n",
    "    return MultResults(results, samples)\n",
    "\n",
    "def get_small_scalar_multiples(mult: MultIdent,\n",
    "                               params: DomainParameters,\n",
    "                               bits: int,\n",
    "                               samples: int = 100,\n",
    "                               use_init: bool = True,\n",
    "                               use_multiply: bool = True,\n",
    "                               seed: bytes | None = None,\n",
    "                               kind: str = \"precomp+necessary\") -> MultResults:\n",
    "    from pyecsca.sca.re.rpa import multiples_computed\n",
    "    import random\n",
    "    \n",
    "    results = []\n",
    "    if seed is not None:\n",
    "        random.seed(seed)\n",
    "\n",
    "    # If no countermeasure is used, we have fully random scalars.\n",
    "    # Otherwise, fix one per chunk.\n",
    "    if mult.countermeasure is None:\n",
    "        scalars = [random.randint(1, 2**bits) for _ in range(samples)]\n",
    "    else:\n",
    "        one = random.randint(1, 2**bits)\n",
    "        scalars = [one for _ in range(samples)]\n",
    "\n",
    "    for scalar in scalars:\n",
    "        # Use a list for less memory usage.\n",
    "        results.append(list(multiples_computed(scalar, params, mult.klass, mult.partial, use_init, use_multiply, kind=kind)))\n",
    "    return MultResults(results, samples)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8c5e9543-8447-4362-b9e2-c896d71f69a9",
   "metadata": {},
   "source": [
    "## Prepare"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4d5c7f10-618f-4612-b594-81d1607b0d1d",
   "metadata": {},
   "outputs": [],
   "source": [
    "category = \"secg\"\n",
    "curve = \"secp256r1\"\n",
    "kind = \"precomp+necessary\"\n",
    "use_init = True\n",
    "use_multiply = True\n",
    "params = get_params(category, curve, \"projective\")\n",
    "num_workers = 20\n",
    "bits = params.order.bit_length()\n",
    "samples = 100\n",
    "selected_mults = all_mults"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3aaf712e-5b97-4390-8dd4-e1db1dfe36a2",
   "metadata": {},
   "source": [
    "## Run\n",
    "Run this cell as many times as you want. It will write chunks into files."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "84359084-4116-436c-92cd-d43fdfeca842",
   "metadata": {},
   "outputs": [],
   "source": [
    "multiples_mults = {}\n",
    "chunk_id = randbytes(4).hex()\n",
    "with TaskExecutor(max_workers=num_workers, mp_context=spawn_context) as pool, enable_spawn(get_small_scalar_multiples) as target:\n",
    "    for mult in selected_mults:\n",
    "        for countermeasure in (None, \"gsr\", \"additive\", \"multiplicative\", \"euclidean\", \"bt\"):\n",
    "            mwc = mult.with_countermeasure(countermeasure)\n",
    "            pool.submit_task(mwc,\n",
    "                             target,\n",
    "                             mwc, params, bits, samples, seed=chunk_id, kind=kind, use_init=use_init, use_multiply=use_multiply)\n",
    "    for mult, future in tqdm(pool.as_completed(), desc=\"Computing small scalar distributions.\", total=len(pool.tasks)):\n",
    "        print(f\"Got {mult}.\")\n",
    "        if error := future.exception():\n",
    "            print(\"Error!\", error)\n",
    "            continue\n",
    "        res = future.result()\n",
    "        if mult not in multiples_mults:\n",
    "            multiples_mults[mult] = res\n",
    "        else:\n",
    "            # Accumulate\n",
    "            multiples_mults[mult].merge(res)\n",
    "\n",
    "    # Handle the enable_spawn trick that messes up class modules.\n",
    "    for k, v in multiples_mults.items():\n",
    "        v.__class__ = MultResults\n",
    "        v.__module__ = \"common\"\n",
    "\n",
    "# Save\n",
    "with open(f\"multiples_{bits}_{'init' if use_init else 'noinit'}_{'mult' if use_multiply else 'nomult'}_chunk{chunk_id}.pickle\",\"wb\") as h:\n",
    "    for mult, res in multiples_mults.items():\n",
    "        pickle.dump((mult, res), h)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b4471a1d-fdc3-4be7-bd61-5ddd22180b41",
   "metadata": {},
   "source": [
    "### Load\n",
    "**Beware**, the following load with try to load all chunks into memory, that will be very large.\n",
    "\n",
    "You probably dont want to run this."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d291832-b0c7-4c3a-9989-22079e4e0f53",
   "metadata": {},
   "outputs": [],
   "source": [
    "multiples_mults = {}\n",
    "for fname in glob.glob(f\"multiples_{bits}_{'init' if use_init else 'noinit'}_{'mult' if use_multiply else 'nomult'}_chunk*.pickle\"):\n",
    "    with open(fname, \"rb\") as f:\n",
    "        while True:\n",
    "            try:\n",
    "                mult, vals = pickle.load(f)\n",
    "                if mult not in multiples_mults:\n",
    "                    multiples_mults[mult] = vals\n",
    "                else:\n",
    "                    multiples_mults[mult].merge(vals)\n",
    "            except EOFError:\n",
    "                break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11b447f2-71ab-417e-a856-1724788cfc91",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}