{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "b7b7430c",
   "metadata": {},
   "source": [
    "# Generic register composition\n",
    "\n",
    "**Download Notebook** - {nb-download}`abstract_construction.ipynb`\n",
    "\n",
    "Larger quantum algorithms are often assembled from interchangeable component boxes that act on quantum registers.\n",
    "\n",
    "- A generic Guppy struct stores the component functions.\n",
    "- Its `compose` method defines how their quantum-register types connect.\n",
    "- A generic quantum-register type may be one qubit array or a struct bundling several quantum registers.\n",
    "- Reusing the middle type makes incompatible connections a compile-time error."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bac7ecb0",
   "metadata": {},
   "outputs": [],
   "source": [
    "from guppylang.decorator import guppy\n",
    "from guppylang.std.angles import angle\n",
    "from guppylang.std.builtins import Function, array, nat, output\n",
    "from guppylang.std.quantum import collect_measurements, cx, h, measure_array, qubit, rz\n",
    "\n",
    "from guppyalgos.utils import qarray, transversal"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "generic-register-types",
   "metadata": {},
   "source": [
    "## Generic quantum-register types\n",
    "\n",
    "A generic type can stand for a complete quantum-register shape. The same higher-order function can therefore accept an operation over an array, a tuple, or a struct containing several quantum registers.\n",
    "\n",
    "- The concrete operation and argument determine `Registers` at the call site.\n",
    "- Reusing `Registers` in both positions ensures that the operation accepts the supplied shape.\n",
    "- Guppy checks this compatibility at compile time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "generic-register-operation",
   "metadata": {},
   "outputs": [],
   "source": [
    "@guppy\n",
    "def apply_register_operation[Registers](\n",
    "    operation: Function[[Registers], None],\n",
    "    qregs: Registers,\n",
    ") -> None:\n",
    "    operation(qregs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "array-register",
   "metadata": {},
   "source": [
    "### A qubit array\n",
    "\n",
    "The most common quantum register is `array[qubit, n]`. Its width $n$ is part of its type and is inferred when the function is used.\n",
    "\n",
    "An $n$-qubit register has state space\n",
    "\n",
    "$$\n",
    "\\mathcal{H}_n = (\\mathbb{C}^2)^{\\otimes n}.\n",
    "$$\n",
    "\n",
    "The operation below applies a Hadamard gate to every qubit, giving $H^{\\otimes n}$. On an all-zero input, this prepares a uniform superposition:\n",
    "\n",
    "$$\n",
    "|0\\rangle^{\\otimes n}\n",
    "\\longmapsto\n",
    "|+\\rangle^{\\otimes n}\n",
    "= \\frac{1}{\\sqrt{2^n}}\\sum_{x=0}^{2^n-1}|x\\rangle.\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "array-register-operation",
   "metadata": {},
   "outputs": [],
   "source": [
    "@guppy\n",
    "def array_operation[n: nat](qreg: array[qubit, n]) -> None:\n",
    "    transversal(h, qreg)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "tuple-register",
   "metadata": {},
   "source": [
    "### A tuple of qubits\n",
    "\n",
    "A fixed tuple is useful when its positions have distinct roles. Here the tuple describes a two-qubit target. Starting from $|00\\rangle$, the following operation prepares a Bell pair:\n",
    "\n",
    "$$\n",
    "|00\\rangle\n",
    "\\xrightarrow{H\\otimes I}\n",
    "\\frac{|00\\rangle+|10\\rangle}{\\sqrt{2}}\n",
    "\\xrightarrow{\\mathrm{CX}}\n",
    "\\frac{|00\\rangle+|11\\rangle}{\\sqrt{2}}.\n",
    "$$\n",
    "\n",
    "In this expression, the first ket position is the first tuple element and controls the CX gate."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "tuple-register-operation",
   "metadata": {},
   "outputs": [],
   "source": [
    "@guppy\n",
    "def tuple_operation(qregs: tuple[qubit, qubit]) -> None:\n",
    "    h(qregs[0])\n",
    "    cx(qregs[0], qregs[1])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "struct-register",
   "metadata": {},
   "source": [
    "### A struct containing quantum registers\n",
    "\n",
    "A Guppy struct gives names to a bundle of related quantum registers. The complete struct can be used as one generic quantum-register type while its operation still accesses each field directly.\n",
    "\n",
    "Here, the work and target registers each contain $n$ qubits. Together, they have state space $\\mathcal{H}_{\\mathrm{work}}\\otimes\\mathcal{H}_{\\mathrm{target}}$. Starting with both registers at zero, the operation creates $n$ Bell pairs:\n",
    "\n",
    "$$\n",
    "|0\\rangle_{\\mathrm{work}}^{\\otimes n}|0\\rangle_{\\mathrm{target}}^{\\otimes n}\n",
    "\\longmapsto\n",
    "\\frac{1}{\\sqrt{2^n}}\\sum_{x=0}^{2^n-1}\n",
    "|x\\rangle_{\\mathrm{work}}|x\\rangle_{\\mathrm{target}}.\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "struct-register-operation",
   "metadata": {},
   "outputs": [],
   "source": [
    "@guppy.struct\n",
    "class WorkAndTarget[n: nat]:\n",
    "    work_qreg: array[qubit, n]\n",
    "    target_qreg: array[qubit, n]\n",
    "\n",
    "\n",
    "@guppy\n",
    "def struct_operation[n: nat](qregs: WorkAndTarget[n]) -> None:\n",
    "    transversal(h, qregs.work_qreg)\n",
    "    transversal(cx, qregs.work_qreg, qregs.target_qreg)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "multiple-generic-registers",
   "metadata": {},
   "source": [
    "### Several generic quantum-register types\n",
    "\n",
    "Composed algorithms can keep quantum registers with different roles independent. For example, `PrepRegisters` selects an operation and `TargetRegisters` is acted on by it. Either type may be an array, tuple, or struct."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "multiple-generic-register-operation",
   "metadata": {},
   "outputs": [],
   "source": [
    "@guppy\n",
    "def apply_select[PrepRegisters, TargetRegisters](\n",
    "    select: Function[[PrepRegisters, TargetRegisters], None],\n",
    "    prep_qreg: PrepRegisters,\n",
    "    target_qreg: TargetRegisters,\n",
    ") -> None:\n",
    "    select(prep_qreg, target_qreg)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "20a83622",
   "metadata": {},
   "source": [
    "## Advanced register composition\n",
    "\n",
    "`BoxComposition` follows the same pattern as composable algorithm structs such as `LCU`: it owns its component boxes and exposes a typed `compose` method.\n",
    "\n",
    "- `box_0` accepts `Register0` and `Register1`.\n",
    "- `box_1` accepts the same `Register1` type and `Register2`.\n",
    "- The boxes may otherwise have different function signatures; only `box_1` accepts an angle.\n",
    "\n",
    "Let $\\mathcal{H}_j$ be the state space of `Registerj`, and let $U_0$ and $U_1(\\theta)$ represent the two boxes. The complete operation acts on $\\mathcal{H}_0\\otimes\\mathcal{H}_1\\otimes\\mathcal{H}_2$:\n",
    "\n",
    "$$\n",
    "U_{\\mathrm{compose}}(\\theta)\n",
    "= \\bigl(I_0\\otimes U_1(\\theta)\\bigr)\n",
    "  \\bigl(U_0\\otimes I_2\\bigr).\n",
    "$$\n",
    "\n",
    "Read the product from right to left: `box_0` acts first on registers $0$ and $1$, then `box_1` acts on registers $1$ and $2$. Each identity $I_j$ leaves the unused register unchanged.\n",
    "\n",
    "The repeated `Register1` type ensures that both boxes agree on the shared register's complete shape."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a6a65f0",
   "metadata": {},
   "outputs": [],
   "source": [
    "@guppy.struct\n",
    "class BoxComposition[Register0, Register1, Register2]:\n",
    "    box_0: Function[[Register0, Register1], None]\n",
    "    box_1: Function[[Register1, Register2, angle], None]\n",
    "\n",
    "    @guppy\n",
    "    def compose(\n",
    "        self,\n",
    "        qreg_0: Register0,\n",
    "        qreg_1: Register1,\n",
    "        qreg_2: Register2,\n",
    "        theta: angle,\n",
    "    ) -> None:\n",
    "        self.box_0(qreg_0, qreg_1)\n",
    "        self.box_1(qreg_1, qreg_2, theta)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "single-bundles",
   "metadata": {},
   "source": [
    "## Every bundle contains one quantum register\n",
    "\n",
    "First, each generic type is instantiated as one `array[qubit, n]`. The three registers therefore contain $3n$ qubits in total:\n",
    "\n",
    "$$\n",
    "\\mathcal{H}_{\\mathrm{total}}\n",
    "= \\mathcal{H}_n\\otimes\\mathcal{H}_n\\otimes\\mathcal{H}_n.\n",
    "$$\n",
    "\n",
    "The component functions are specialized before they are stored because higher-rank polymorphic function values are not supported."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a40e5727",
   "metadata": {},
   "outputs": [],
   "source": [
    "@guppy\n",
    "def bell_transversal[n: nat](\n",
    "    left_qreg: array[qubit, n], shared: array[qubit, n]\n",
    ") -> None:\n",
    "    transversal(h, left_qreg)\n",
    "    transversal(cx, left_qreg, shared)\n",
    "\n",
    "\n",
    "@guppy\n",
    "def rz_cx_transversal[n: nat](\n",
    "    shared: array[qubit, n], right_qreg: array[qubit, n], theta: angle\n",
    ") -> None:\n",
    "    for i in range(len(shared)):\n",
    "        rz(shared[i], theta)\n",
    "    transversal(cx, shared, right_qreg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "single-main",
   "metadata": {},
   "outputs": [],
   "source": [
    "n_qubits = 3\n",
    "\n",
    "\n",
    "@guppy\n",
    "def single_register_bundles() -> None:\n",
    "    qreg_0 = qarray(n_qubits)\n",
    "    qreg_1 = qarray(n_qubits)\n",
    "    qreg_2 = qarray(n_qubits)\n",
    "\n",
    "    composition = BoxComposition(\n",
    "        bell_transversal[n_qubits],\n",
    "        rz_cx_transversal[n_qubits],\n",
    "    )\n",
    "    composition.compose(\n",
    "        qreg_0, qreg_1, qreg_2, angle(0.25)\n",
    "    )\n",
    "\n",
    "    output(\"qreg_0\", collect_measurements(measure_array(qreg_0)))\n",
    "    output(\"qreg_1\", collect_measurements(measure_array(qreg_1)))\n",
    "    output(\"qreg_2\", collect_measurements(measure_array(qreg_2)))\n",
    "\n",
    "\n",
    "single_results = (\n",
    "    single_register_bundles.emulator(n_qubits=3 * n_qubits)\n",
    "    .with_seed(42)\n",
    "    .with_shots(100)\n",
    "    .run()\n",
    ")\n",
    "print(\"Single-register bundles:\", single_results.collated_counts())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "middle-bundle",
   "metadata": {},
   "source": [
    "## The middle bundle contains two quantum registers\n",
    "\n",
    "The composer itself does not change. Instead, `Register1` is instantiated as `MiddleBundle[n]`, containing two independently addressable quantum registers:\n",
    "\n",
    "$$\n",
    "\\mathcal{H}_1\n",
    "= \\mathcal{H}_{\\mathrm{upper}}\\otimes\\mathcal{H}_{\\mathrm{lower}}\n",
    "= \\mathcal{H}_n\\otimes\\mathcal{H}_n.\n",
    "$$\n",
    "\n",
    "The middle bundle now contains $2n$ qubits, bringing the total to $4n$.\n",
    "\n",
    "The composition equation above still applies: only the shape of $\\mathcal{H}_1$ has changed. Both boxes must accept the complete bundle, so the shared boundary remains type safe."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "middle-struct",
   "metadata": {},
   "outputs": [],
   "source": [
    "@guppy.struct\n",
    "class MiddleBundle[n: nat]:\n",
    "    upper_qreg: array[qubit, n]\n",
    "    lower_qreg: array[qubit, n]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "middle-boxes",
   "metadata": {},
   "outputs": [],
   "source": [
    "@guppy\n",
    "def bundled_box_0[n: nat](\n",
    "    left_qreg: array[qubit, n], middle: MiddleBundle[n]\n",
    ") -> None:\n",
    "    transversal(h, left_qreg)\n",
    "    transversal(cx, left_qreg, middle.upper_qreg)\n",
    "    transversal(cx, left_qreg, middle.lower_qreg)\n",
    "\n",
    "\n",
    "@guppy\n",
    "def bundled_box_1[n: nat](\n",
    "    middle: MiddleBundle[n], right_qreg: array[qubit, n], theta: angle\n",
    ") -> None:\n",
    "    for i in range(n):\n",
    "        rz(middle.upper_qreg[i], theta)\n",
    "        rz(middle.lower_qreg[i], theta)\n",
    "    transversal(cx, middle.upper_qreg, right_qreg)\n",
    "    transversal(cx, middle.lower_qreg, right_qreg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "middle-main",
   "metadata": {},
   "outputs": [],
   "source": [
    "@guppy\n",
    "def two_register_middle_bundle() -> None:\n",
    "    qreg_0 = qarray(n_qubits)\n",
    "    middle = MiddleBundle(qarray(n_qubits), qarray(n_qubits))\n",
    "    qreg_2 = qarray(n_qubits)\n",
    "\n",
    "    composition = BoxComposition(\n",
    "        bundled_box_0[n_qubits],\n",
    "        bundled_box_1[n_qubits],\n",
    "    )\n",
    "    composition.compose(\n",
    "        qreg_0, middle, qreg_2, angle(0.25)\n",
    "    )\n",
    "\n",
    "    output(\"qreg_0\", collect_measurements(measure_array(qreg_0)))\n",
    "    output(\"middle_upper\", collect_measurements(measure_array(middle.upper_qreg)))\n",
    "    output(\"middle_lower\", collect_measurements(measure_array(middle.lower_qreg)))\n",
    "    output(\"qreg_2\", collect_measurements(measure_array(qreg_2)))\n",
    "\n",
    "\n",
    "bundled_results = (\n",
    "    two_register_middle_bundle.emulator(n_qubits=4 * n_qubits)\n",
    "    .with_seed(42)\n",
    "    .with_shots(100)\n",
    "    .run()\n",
    ")\n",
    "print(\"Two-register middle bundle:\", bundled_results.collated_counts())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0d08357f",
   "metadata": {},
   "source": [
    "## Development practices\n",
    "\n",
    "- Store related component functions in a generic Guppy struct when they form one reusable abstraction.\n",
    "- Put the wiring in the struct's `compose` method.\n",
    "- Use `Function`, rather than `typing.Callable`, for Guppy function values.\n",
    "- Reuse the same type variable at connected boundaries so Guppy checks compatibility.\n",
    "- A type parameter can represent an array or a struct containing several quantum registers."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "guppyalgos (3.13.3.final.0)",
   "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.13.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}