{
    "cells": [
        {
            "cell_type": "markdown",
            "id": "cb1df4a3",
            "metadata": {},
            "source": [
                "# QROM with unary iteration\n",
                "\n",
                "**Download Notebook** - {nb-download}`qrom_unary_iteration.ipynb`\n",
                "\n",
                "- Unary iteration QROM implementation.\n",
                "\n",
                "- QROM is used to store classical data in a quantum circuit. It takes as input an index register and a data register, and for each possible value of the index register, it writes the corresponding data value to the data register."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 1,
            "id": "86339f20",
            "metadata": {},
            "outputs": [],
            "source": [
                "from guppyalgos.algorithms.select.qrom import qrom_unary_iteration\n",
                "\n",
                "from guppylang.decorator import guppy\n",
                "from guppylang.std.builtins import array, output\n",
                "from guppylang.std.quantum import measure_array, qubit, toffoli, x, h, collect_measurements\n",
                "\n",
                "from guppyalgos.utils import int_to_bits, qarray\n",
                "\n",
                "from guppylang import comptime\n",
                "\n",
                "from guppyalgos.primitives.gate_decompositions.and_op import (\n    temp_and_compute,\n    temp_and_uncompute,\n)"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "5ac361aa",
            "metadata": {},
            "source": [
                "- First we must build the input data\n",
                "- A list of lists of bools is used to represent the classical data to be stored in the QROM.\n",
                "- Currently only a single target register is supported."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 2,
            "id": "422c46ef",
            "metadata": {},
            "outputs": [],
            "source": [
                "import numpy as np\n",
                "# build a random of list of list of bools\n",
                "\n",
                "n_index_qubits = 4\n",
                "n_state_qubits = 4\n",
                "n_index_elements = 2**n_index_qubits  # check when this is not a power of 2\n",
                "\n",
                "data_input = [np.random.choice([False, True], n_state_qubits).tolist() for _ in range(n_index_elements)]"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "90ba38b8",
            "metadata": {},
            "source": [
                "- We use a metaprogramming approach to build the QROM circuit where we use a python function to generate the guppy function\n",
                "- The QROM is has freedom to use different compute and uncompute functions for the AND operations\n",
                "\n",
                "- In this example function we initialize the index register to a specific value and check that the correct data is written to the data register by measuring the index register at the end of the circuit, which should still be in the same index."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 3,
            "id": "4a42a06a",
            "metadata": {},
            "outputs": [],
            "source": [
                "\n",
                "def qrom_and_input(compute_and, uncompute_and, data_input, index):\n",
                "\n",
                "    binary_index = int_to_bits(index, n_index_qubits)\n",
                "\n",
                "    qrom = qrom_unary_iteration(data_input, comp_and_op=compute_and, uncomp_and_op=uncompute_and)\n",
                "\n",
                "    @guppy\n",
                "    def main() -> None:\n",
                "        idx = binary_index\n",
                "        index_qreg = qarray(n_index_qubits)\n",
                "        state_qreg = qarray(n_state_qubits)\n",
                "\n",
                "        for bit in range(n_index_qubits):\n",
                "            if idx[bit]:\n",
                "                x(index_qreg[bit])\n",
                "\n",
                "        qrom(index_qreg, state_qreg)\n",
                "\n",
                "        output(\"index\", collect_measurements(measure_array(index_qreg)))\n",
                "        output(\"state\", collect_measurements(measure_array(state_qreg)))\n",
                "\n",
                "\n",
                "    total_qubits = n_state_qubits + n_index_qubits + (n_index_qubits - 1)\n",
                "    my_shots = main.emulator(n_qubits=total_qubits).with_seed(42).with_shots(1).run()\n",
                "\n",
                "    print(f\"Measured data {index}: \", my_shots.collated_shots()[0][\"state\"])\n",
                "    print(f\"Expected data {index}: \", data_input[index])"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 4,
            "id": "b92420ff",
            "metadata": {},
            "outputs": [
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "--- Test 0 ---\n",
                        "Measured data 0:  [[0, 0, 1, 1]]\n",
                        "Expected data 0:  [False, False, True, True]\n",
                        "--- Test 1 ---\n",
                        "Measured data 1:  [[1, 0, 1, 0]]\n",
                        "Expected data 1:  [True, False, True, False]\n",
                        "--- Test 2 ---\n",
                        "Measured data 2:  [[1, 0, 1, 1]]\n",
                        "Expected data 2:  [True, False, True, True]\n",
                        "--- Test 3 ---\n",
                        "Measured data 3:  [[0, 0, 0, 1]]\n",
                        "Expected data 3:  [False, False, False, True]\n"
                    ]
                }
            ],
            "source": [
                "for i in range(4):\n",
                "    print(f\"--- Test {i} ---\")\n",
                "    qrom_and_input(toffoli, toffoli, data_input, i)"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "add3a8b9",
            "metadata": {},
            "source": [
                "- We can also use T gate efficient AND functions to reduce the T gate count of the QROM circuit.\n",
                "- Following the implementation from https://arxiv.org/abs/1805.03662 we can do the compute temporary and with 4 T gates and the uncompute with 0 T gates using measurement based computation."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 5,
            "id": "17a68111",
            "metadata": {},
            "outputs": [
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "Measured data 3:  [[0, 0, 0, 1]]\n",
                        "Expected data 3:  [False, False, False, True]\n"
                    ]
                }
            ],
            "source": [
                "qrom_and_input(temp_and_compute, temp_and_uncompute, data_input, 3)"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "13556a92",
            "metadata": {},
            "source": [
                "- Here we can apply a uniform superposition to the index register by applying Hadamard gates to all qubits in the index register.\n",
                "- Each shot will give a uniformly random binary index and the corresponding data value.\n"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 6,
            "id": "c7b8c584",
            "metadata": {},
            "outputs": [],
            "source": [
                "from guppyalgos.utils import transversal\n",
                "\n",
                "def qrom_h_transversal(compute_and, uncompute_and, data_input):\n",
                "\n",
                "    qrom = qrom_unary_iteration(data_input, comp_and_op=compute_and, uncomp_and_op=uncompute_and)\n",
                "\n",
                "    @guppy\n",
                "    def main() -> None:\n",
                "        index_qreg = qarray(n_index_qubits)\n",
                "        state_qreg = qarray(n_state_qubits)\n",
                "\n",
                "        transversal(h, index_qreg)\n",
                "\n",
                "        qrom(index_qreg, state_qreg)\n",
                "\n",
                "        output(\"index\", collect_measurements(measure_array(index_qreg)))\n",
                "        output(\"state\", collect_measurements(measure_array(state_qreg)))\n",
                "\n",
                "\n",
                "    total_qubits = n_state_qubits + n_index_qubits + (n_index_qubits - 1)\n",
                "    my_shots = main.emulator(n_qubits=total_qubits).with_seed(42).with_shots(10).run()\n",
                "\n",
                "    print(f\"Measured data: \", my_shots.collated_shots())"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 7,
            "id": "d4e5555c",
            "metadata": {},
            "outputs": [
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "Measured data:  [{'index': [[1, 0, 0, 1]], 'state': [[0, 0, 0, 1]]}, {'index': [[0, 0, 1, 1]], 'state': [[0, 0, 0, 1]]}, {'index': [[0, 1, 1, 1]], 'state': [[1, 1, 1, 0]]}, {'index': [[0, 0, 0, 0]], 'state': [[0, 0, 1, 1]]}, {'index': [[1, 0, 0, 0]], 'state': [[0, 1, 0, 0]]}, {'index': [[0, 0, 0, 1]], 'state': [[1, 0, 1, 0]]}, {'index': [[1, 0, 0, 1]], 'state': [[0, 0, 0, 1]]}, {'index': [[0, 0, 0, 1]], 'state': [[1, 0, 1, 0]]}, {'index': [[0, 1, 0, 0]], 'state': [[0, 1, 0, 0]]}, {'index': [[0, 1, 1, 1]], 'state': [[1, 1, 1, 0]]}]\n"
                    ]
                }
            ],
            "source": [
                "qrom_h_transversal(toffoli, toffoli, data_input)"
            ]
        }
    ],
    "metadata": {
        "kernelspec": {
            "display_name": "guppyalgos (3.13.3)",
            "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
}