{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "241f8419",
   "metadata": {},
   "source": [
    "# Modular Multiplier Example\n",
    "\n",
    "**Download Notebook** - {nb-download}`multiplier.ipynb`\n",
    "\n",
    "This notebook demonstrates the Gidney ripple-based multiplier circuits in `guppyalgos.primitives.arithmetic`:\n",
    "\n",
    "- `multiplier_ripple_gidney_mod`\n",
    "- `cntrl_multiplier_ripple_gidney_mod`\n",
    "\n",
    "Both implement multiplication modulo $2^n$ on little-endian bit registers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "8353c631",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import no_type_check\n",
    "\n",
    "from guppylang import guppy\n",
    "from guppylang.std.builtins import array, output\n",
    "from guppylang.std.quantum import (\n",
    "    collect_measurements,\n",
    "    discard,\n",
    "    measure_array,\n",
    "    qubit,\n",
    "    x,\n",
    " )\n",
    "\n",
    "from guppyalgos.primitives.arithmetic import (\n",
    "    cntrl_multiplier_ripple_gidney_mod,\n",
    "    multiplier_ripple_gidney_mod,\n",
    " )\n",
    "from guppyalgos.utils import apply_bitstring, bits_to_int, int_to_bits, qarray"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f19615c4",
   "metadata": {},
   "source": [
    "## Math and Circuit Structure\n",
    "\n",
    "For $n$-bit little-endian registers $a$ and $b$, the target operation is\n",
    "\n",
    "$$\n",
    "|a\\rangle|b\\rangle|p\\rangle \\mapsto |a\\rangle|b\\rangle|p + a\\cdot b \\; (\\mathrm{mod}\\; 2^n)\\rangle.\n",
    "$$\n",
    "\n",
    "Write $b = \\sum_{i=0}^{n-1} b_i 2^i$. Then\n",
    "\n",
    "$$\n",
    "a\\cdot b = \\sum_{i=0}^{n-1} b_i\\,(a\\ll i).\n",
    "$$\n",
    "\n",
    "So the circuit performs a sequence of modular additions of shifted addends.\n",
    "\n",
    "- Uncontrolled multiplier: add $(a\\ll i)$ into `product` only when multiplier bit $b_i=1$.\n",
    "- Controlled multiplier: same, but each addition is additionally gated by a global control `ctrl`.\n",
    "\n",
    "Internally, temporary work qubits are created and uncomputed each step, then discarded. The total qubit count for the multiplier breaks down into:\n",
    "- $n$ qubits for the input $a$\n",
    "- $n$ qubits for the input $b$\n",
    "- $n$ qubits for the (temporary) partial product $b_i\\,(a\\ll i)$\n",
    "- $n$ qubits for the output $a \\cdot b$\n",
    "- Any ancilla qubits used by the adder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "2667d5d6",
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_multiplier_case(n: int, a: int, b: int) -> dict:\n",
    "    \"\"\"Run one uncontrolled modular multiplication case and return measured bits.\"\"\"\n",
    "    a_bits = int_to_bits(a, n)\n",
    "    b_bits = int_to_bits(b, n)\n",
    "    a_bit_array = array(*a_bits)\n",
    "    b_bit_array = array(*b_bits)\n",
    "\n",
    "    @guppy\n",
    "    @no_type_check\n",
    "    def main() -> None:\n",
    "        a_reg = qarray(n)\n",
    "        apply_bitstring(a_reg, a_bit_array)\n",
    "        multiplier = qarray(n)\n",
    "        apply_bitstring(multiplier, b_bit_array)\n",
    "        product = qarray(n)\n",
    "\n",
    "        multiplier_ripple_gidney_mod(a_reg, multiplier, product)\n",
    "\n",
    "        output(\"a_meas\", collect_measurements(measure_array(a_reg)))\n",
    "        output(\"mult_meas\", collect_measurements(measure_array(multiplier)))\n",
    "        output(\"prod_meas\", collect_measurements(measure_array(product)))\n",
    "\n",
    "    res = main.emulator(n_qubits=5 * n - 1).run()\n",
    "    return res.results[0].as_dict()\n",
    "\n",
    "\n",
    "def run_cntrl_multiplier_case(\n",
    "    n: int, a: int, b: int, ctrl_active: bool\n",
    " ) -> dict:\n",
    "    \"\"\"Run one controlled modular multiplication case and return measured bits.\"\"\"\n",
    "    a_bits = int_to_bits(a, n)\n",
    "    b_bits = int_to_bits(b, n)\n",
    "    a_bit_array = array(*a_bits)\n",
    "    b_bit_array = array(*b_bits)\n",
    "\n",
    "    @guppy\n",
    "    @no_type_check\n",
    "    def main() -> None:\n",
    "        a_reg = qarray(n)\n",
    "        apply_bitstring(a_reg, a_bit_array)\n",
    "        multiplier = qarray(n)\n",
    "        apply_bitstring(multiplier, b_bit_array)\n",
    "        product = qarray(n)\n",
    "\n",
    "        ctrl = qubit()\n",
    "        if ctrl_active:\n",
    "            x(ctrl)\n",
    "\n",
    "        cntrl_multiplier_ripple_gidney_mod(ctrl, a_reg, multiplier, product)\n",
    "\n",
    "        output(\"a_meas\", collect_measurements(measure_array(a_reg)))\n",
    "        output(\"mult_meas\", collect_measurements(measure_array(multiplier)))\n",
    "        output(\"prod_meas\", collect_measurements(measure_array(product)))\n",
    "\n",
    "        discard(ctrl)\n",
    "\n",
    "    res = main.emulator(n_qubits=5 * n + 1).run()\n",
    "    return res.results[0].as_dict()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "131f7c5b",
   "metadata": {},
   "source": [
    "## Uncontrolled Multiplier Demo\n",
    "\n",
    "The following cell checks that\n",
    "\n",
    "$$\n",
    "\\mathrm{product}_{\\mathrm{out}} = (a\\cdot b) \\bmod 2^n\n",
    "$$\n",
    "\n",
    "and that the input registers are preserved."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "f624b999",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "n=2, a=1, b=1 -> prod=1 (expected 1)\n",
      "n=2, a=2, b=3 -> prod=2 (expected 2)\n",
      "n=3, a=3, b=2 -> prod=6 (expected 6)\n",
      "n=4, a=5, b=7 -> prod=3 (expected 3)\n",
      "Uncontrolled multiplier checks passed.\n"
     ]
    }
   ],
   "source": [
    "uncontrolled_cases = [(2, 1, 1), (2, 2, 3), (3, 3, 2), (4, 5, 7)]\n",
    "\n",
    "for n, a, b in uncontrolled_cases:\n",
    "    result = run_multiplier_case(n, a, b)\n",
    "    expected_bits = int_to_bits((a * b) % (2**n), n)\n",
    "\n",
    "    assert result[\"a_meas\"] == int_to_bits(a, n)\n",
    "    assert result[\"mult_meas\"] == int_to_bits(b, n)\n",
    "    assert result[\"prod_meas\"] == expected_bits\n",
    "\n",
    "    print(\n",
    "        f\"n={n}, a={a}, b={b} -> prod={bits_to_int(result['prod_meas'])} (expected {bits_to_int(expected_bits)})\"\n",
    "    )\n",
    "\n",
    "print(\"Uncontrolled multiplier checks passed.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d3c1ebd3",
   "metadata": {},
   "source": [
    "## Controlled Multiplier Demo\n",
    "\n",
    "For the controlled circuit, we check two behaviors:\n",
    "\n",
    "- `ctrl = 0`: product register is unchanged.\n",
    "- `ctrl = 1`: product register becomes $(a\\cdot b) \\bmod 2^n$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "14143e5b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ctrl=0, n=2, a=1, b=1 ->  prod=0 (expected 0)\n",
      "ctrl=1, n=2, a=1, b=1 ->  prod=1 (expected 1)\n",
      "ctrl=0, n=2, a=2, b=3 ->  prod=0 (expected 0)\n",
      "ctrl=1, n=2, a=2, b=3 ->  prod=2 (expected 2)\n",
      "ctrl=0, n=3, a=3, b=2 ->  prod=0 (expected 0)\n",
      "ctrl=1, n=3, a=3, b=2 ->  prod=6 (expected 6)\n",
      "ctrl=0, n=4, a=5, b=7 ->  prod=0 (expected 0)\n",
      "ctrl=1, n=4, a=5, b=7 ->  prod=3 (expected 3)\n",
      "Controlled multiplier checks passed.\n"
     ]
    }
   ],
   "source": [
    "controlled_cases = [(2, 1, 1), (2, 2, 3), (3, 3, 2), (4, 5, 7)]\n",
    "\n",
    "for n, a, b in controlled_cases:\n",
    "    for ctrl_active in (False, True):\n",
    "        result = run_cntrl_multiplier_case(n, a, b, ctrl_active)\n",
    "        expected_value = 0 if not ctrl_active else (a * b) % (2**n)\n",
    "        expected_bits = int_to_bits(expected_value, n)\n",
    "\n",
    "        assert result[\"a_meas\"] == int_to_bits(a, n)\n",
    "        assert result[\"mult_meas\"] == int_to_bits(b, n)\n",
    "        assert result[\"prod_meas\"] == expected_bits\n",
    "\n",
    "        print(\n",
    "            f\"ctrl={int(ctrl_active)}, n={n}, a={a}, b={b} -> \",\n",
    "            f\"prod={bits_to_int(result['prod_meas'])} (expected {bits_to_int(expected_bits)})\",\n",
    "        )\n",
    "\n",
    "print(\"Controlled multiplier checks passed.\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "guppyalgos (3.12.x)",
   "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.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}