{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "391af084",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pennylane as pl\n",
    "from pennylane import numpy as np\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d6c47b4",
   "metadata": {},
   "source": [
    "# Exercise 1: Uniform superposition"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3858c4b2",
   "metadata": {},
   "outputs": [],
   "source": [
    "n_bits = 2\n",
    "dev = pl.device(\"default.qubit\", wires=n_bits)\n",
    "wires = list(range(n_bits))\n",
    "\n",
    "@pl.qnode(dev)\n",
    "def circuit():\n",
    "    pl.Snapshot(\"Initial state\")\n",
    "    # build a uniform superposition\n",
    "    pl.Snapshot(\"After applying the Hadamard gates\")\n",
    "    return pl.probs(wires=wires)  \n",
    "\n",
    "pl.drawer.use_style(\"black_white\")\n",
    "pl.draw_mpl(circuit)();\n",
    "results = pl.snapshots(circuit)()\n",
    "for k, result in results.items():\n",
    "    print(f\"{k}: {result}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bafbbeba",
   "metadata": {},
   "outputs": [],
   "source": [
    "y = np.real(results[\"After applying the Hadamard gates\"])\n",
    "bit_strings = [f\"{x:0{n_bits}b}\" for x in range(len(y))]\n",
    "plt.bar(bit_strings, y, color = \"#70CEFF\")\n",
    "plt.xticks(rotation=\"vertical\")\n",
    "plt.xlabel(\"State label\")\n",
    "plt.ylabel(\"Probability Amplitude\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "83f70d12",
   "metadata": {},
   "source": [
    "# Exercise 2: Oracle as a reflection\n",
    "\n",
    "Recall that for one marked key $w$, Grover's phase oracle is\n",
    "\n",
    "$$O_w=I-2|w\\rangle\\langle w|.$$\n",
    "\n",
    "This is a **reflection**: it changes the sign of the marked state's amplitude and leaves the other computational-basis amplitudes unchanged.\n",
    "\n",
    "We will construct it as a quantum circuit in **two steps**:\n",
    "\n",
    "1. **Check the key.** Construct the reversible oracle $U_f|x\\rangle|b\\rangle=|x\\rangle|b\\oplus f(x)\\rangle$, where $f(x)=1$ precisely when $x=w$. To check a key such as `01`, apply $X$ to the wires where the key contains a `0`, use a multi-controlled $X$ on the extra qubit, and undo the initial $X$ gates.\n",
    "2. **Turn the check into a phase flip.** Prepare the extra qubit in $|-\\rangle=(|0\\rangle-|1\\rangle)/\\sqrt2$. Since $X|-\\rangle=-|-\\rangle$, the same checker now acts as $O_w$ on the data qubits:\n",
    "\n",
    "$$U_f(|x\\rangle|-\\rangle)=(-1)^{f(x)}|x\\rangle|-\\rangle.$$\n",
    "\n",
    "For this exercise, mark just **one key**, `01`. You can use `pl.ctrl(pl.PauliX, control=wires)(wires=flag_wire)` for the multi-controlled $X$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bd14604c",
   "metadata": {},
   "outputs": [],
   "source": [
    "keys = [[0,1]]  # one marked key\n",
    "flag_wire = n_bits\n",
    "dev = pl.device(\"default.qubit\", wires=n_bits + 1)\n",
    "\n",
    "def oracle(keys):\n",
    "    # build the reversible checker for keys[0]\n",
    "    # apply X to the data wires where the marked bit is 0\n",
    "    # apply a controlled X to the extra (flag) qubit\n",
    "    # undo the first X gates\n",
    "    pass"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d7087e9d",
   "metadata": {},
   "source": [
    "Now apply the oracle to a uniform superposition. Prepare the extra qubit in $|-\\rangle$ using $X$ followed by $H$, and use the snapshots to observe the **reflection**: the amplitude of $|01\\rangle$ becomes negative, but its measurement probability does not change."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "439bcbb1",
   "metadata": {},
   "outputs": [],
   "source": [
    "@pl.qnode(dev)\n",
    "def circuit(keys):\n",
    "    for i in range(n_bits):\n",
    "        pl.Hadamard(wires=i)\n",
    "    # prepare the extra qubit in |->\n",
    "    pl.Snapshot(\"Before oracle\")\n",
    "    # apply the oracle\n",
    "    pl.Snapshot(\"After oracle\")\n",
    "    return pl.state()\n",
    "\n",
    "results = pl.snapshots(circuit)(keys)\n",
    "\n",
    "# The flag stays in |->. Display the data amplitudes using its |0> component.\n",
    "y1 = np.sqrt(2) * np.real(results[\"Before oracle\"][::2])\n",
    "y2 = np.sqrt(2) * np.real(results[\"After oracle\"][::2])\n",
    "bit_strings = [f\"{x:0{n_bits}b}\" for x in range(len(y1))]\n",
    "pos = np.arange(len(y1))\n",
    "plt.bar(pos - 0.2, y1, width=0.4, color=\"#70CEFF\", label=\"Before oracle\")\n",
    "plt.bar(pos + 0.2, y2, width=0.4, color=\"#C756B2\", label=\"After oracle\")\n",
    "plt.xticks(pos, bit_strings)\n",
    "plt.axhline(0, color=\"black\", linewidth=0.8)\n",
    "plt.xlabel(\"State label\")\n",
    "plt.ylabel(\"Probability amplitude\")\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2b008098",
   "metadata": {},
   "source": [
    "# Exercise 3: Diffusion operator\n",
    "\n",
    "The second reflection in Grover's algorithm is the **diffusion operator**:\n",
    "\n",
    "$$D=2|s\\rangle\\langle s|-I,\\qquad |s\\rangle=H^{\\otimes n}|0^n\\rangle.$$\n",
    "\n",
    "Unlike the oracle's reflection, this one reflects the amplitudes **about their mean**. Together, the two reflections amplify the marked state's amplitude.\n",
    "\n",
    "To implement diffusion, apply $H$ to each data qubit, then $X$ to each data qubit, then a multi-controlled $Z$, and finally undo the $X$ and $H$ gates in reverse order. Use\n",
    "\n",
    "```python\n",
    "pl.ctrl(pl.PauliZ, control=wires[:-1])(wires=wires[-1])\n",
    "```\n",
    "\n",
    "This implementation gives $-D$ rather than $D$, which differs only by an overall global phase. The extra flag qubit is **not** part of diffusion."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a10b0f2e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def diffusion_operator(wires):\n",
    "    # build the diffusion operator circuit\n",
    "    pass"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a050822b",
   "metadata": {},
   "source": [
    "# Exercise 4: One Grover iteration\n",
    "\n",
    "A Grover iteration consists of the oracle reflection followed by the diffusion reflection. With two data qubits and one marked key, **one iteration** is enough to reach the marked state with probability $1$.\n",
    "\n",
    "Run the circuit and compare the amplitudes before the oracle, after the oracle, and after diffusion. In particular, check that **the oracle alone does not change the measurement probabilities**."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c02ee540",
   "metadata": {},
   "outputs": [],
   "source": [
    "@pl.qnode(dev)\n",
    "def circuit():\n",
    "    for i in range(n_bits):\n",
    "        pl.Hadamard(wires=i)\n",
    "    pl.PauliX(wires=flag_wire)\n",
    "    pl.Hadamard(wires=flag_wire)\n",
    "    pl.Snapshot(\"Uniform superposition\")\n",
    "    oracle(keys)\n",
    "    pl.Snapshot(\"State marked by oracle\")\n",
    "    diffusion_operator(wires)\n",
    "    pl.Snapshot(\"Amplitude after diffusion\")\n",
    "    return pl.probs(wires=wires)\n",
    "\n",
    "results = pl.snapshots(circuit)()\n",
    "# inspect the results; plot the amplitudes after the oracle and after diffusion"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e9a3e5b3",
   "metadata": {},
   "source": [
    "# Exercise 5: Grover's algorithm\n",
    "\n",
    "Now use **four data qubits** and search for the single marked key $0101$. Build Grover's algorithm from the oracle and diffusion functions above.\n",
    "\n",
    "For $N=2^n$ possible inputs and $M=1$ marked key, let $\\theta=\\arcsin\\sqrt{M/N}$. A suitable number of iterations is\n",
    "\n",
    "$$r_{\\mathrm{opt}}\\approx\\operatorname{round}\\left(\\frac{\\pi}{4\\theta}-\\frac12\\right).$$\n",
    "\n",
    "**Step 5a.** Set up the new circuit and calculate how many iterations to use."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b680e9bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "n_bits = 4\n",
    "wires = list(range(n_bits))\n",
    "flag_wire = n_bits\n",
    "dev = pl.device(\"default.qubit\", wires=n_bits + 1)\n",
    "keys = [[0,1,0,1]]\n",
    "N = 2**n_bits\n",
    "M = 1\n",
    "theta = np.arcsin(np.sqrt(M / N))\n",
    "r_opt = int(np.round(np.pi / (4 * theta) - 0.5))\n",
    "print(\"Suggested Grover iterations:\", r_opt)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9b4d4285",
   "metadata": {},
   "source": [
    "**Step 5b.** Complete the circuit: prepare a uniform superposition on the data qubits and $|-\\rangle$ on the extra qubit, then repeat the oracle and diffusion operations `iterations` times. Measure only the data qubits."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ce208b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "@pl.qnode(dev)\n",
    "def circuit(iterations):\n",
    "    # build the grover circuit, using the oracle and diffusion operator above\n",
    "    # iterate it \"iterations\" times\n",
    "    return pl.probs(wires=wires)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3091315b",
   "metadata": {},
   "source": [
    "**Inspect the result.** Try one and two iterations, then the suggested number $r_{\\mathrm{opt}}$. Plot the final distribution and check whether the most likely outcome is the marked key. You do **not** need to plot a success-probability curve over many iterations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c24f388d",
   "metadata": {},
   "outputs": [],
   "source": [
    "bit_strings = [f\"{x:0{n_bits}b}\" for x in range(2**n_bits)]\n",
    "marked_index = int(\"\".join(map(str, keys[0])), 2)\n",
    "\n",
    "for iterations in [1, 2, r_opt]:\n",
    "    results = circuit(iterations)\n",
    "    print(f\"After {iterations} iteration(s): P(0101) = {results[marked_index]:.4f}\")\n",
    "\n",
    "results = circuit(r_opt)\n",
    "plt.figure(figsize=(8, 3.5))\n",
    "plt.bar(bit_strings, results, color=\"#70CEFF\")\n",
    "plt.xticks(rotation=\"vertical\")\n",
    "plt.xlabel(\"State label\")\n",
    "plt.ylabel(\"Probability\")\n",
    "plt.title(f\"Grover after {r_opt} iterations\")\n",
    "plt.show()\n",
    "print(\"Most likely outcome:\", bit_strings[int(np.argmax(results))])\n",
    "# draw the final circuit"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pennylane (3.14.7)",
   "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.14.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}