{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "ff7370fb",
   "metadata": {},
   "source": [
    "# T state distillation and the `Option` type\n",
    "\n",
    "**Download this notebook - {nb-download}`t_factory.ipynb`**\n",
    "\n",
    "In this example we will demonstrate how to create a T state using magic state distillation, including the use of the `Option` type to indicate success or failure in a repeat-until-success algorithm."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "73e8e28c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "from guppylang.decorator import guppy\n",
    "from guppylang.std.angles import angle, pi\n",
    "from guppylang.std.builtins import array, owned, comptime\n",
    "from guppylang.std.option import Option, nothing, some\n",
    "from guppylang.std.debug import state_result\n",
    "from guppylang.std.quantum import (\n",
    "    cz,\n",
    "    discard,\n",
    "    h,\n",
    "    measure,\n",
    "    qubit,\n",
    "    ry,\n",
    "    rz,\n",
    ")\n",
    "\n",
    "from selene_sim import build, Quest\n",
    "\n",
    "from hugr.qsystem.result import QsysResult\n",
    "\n",
    "np.set_printoptions(precision=4, suppress=True, linewidth=120)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5472c92a",
   "metadata": {},
   "source": [
    "## Preparation\n",
    "\n",
    "We begin by implementing a function to prepare an approximate T state, taken from https://arxiv.org/abs/2310.12106."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "d9646c0f",
   "metadata": {},
   "outputs": [],
   "source": [
    "phi = np.arccos(1 / 3)\n",
    "\n",
    "\n",
    "@guppy\n",
    "def prepare_approx() -> qubit:\n",
    "    q = qubit()\n",
    "    ry(q, angle(comptime(phi)))\n",
    "    rz(q, pi / 4)\n",
    "    return q"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b3acda45",
   "metadata": {},
   "source": [
    "## Distillation\n",
    "\n",
    "This is the inverse of the $[[5,3,1]]$ encoder in figure 3 of https://arxiv.org/abs/2208.01863. \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "704030e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "@guppy\n",
    "def distill(\n",
    "    target: qubit @ owned,\n",
    "    qs: array[qubit, 4] @ owned,\n",
    ") -> tuple[qubit, bool]:\n",
    "    \"\"\"First argument is the target qubit which will be returned from the circuit.\n",
    "    Other arguments are ancillae, which should also be in an approximate T state.\n",
    "    Returns target qubit and a bool, which is true if the distillation succeeded.\n",
    "    \"\"\"\n",
    "    cz(qs[0], qs[1])\n",
    "    cz(qs[2], qs[3])\n",
    "    cz(target, qs[0])\n",
    "    cz(qs[1], qs[2])\n",
    "    cz(target, qs[3])\n",
    "    # Measuring gives false for success, true for failure.\n",
    "    # We check for all falses to say whether distillation succeeded.\n",
    "    for i in range(4):\n",
    "        h(qs[i])\n",
    "    bits = array(not measure(q) for q in qs)\n",
    "    # Guppy doesn't yet support the `any` or `all` operators yet.\n",
    "    success = True\n",
    "    for b in bits:\n",
    "        success &= b\n",
    "    return target, success"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2ff9a83c",
   "metadata": {},
   "source": [
    "## Repeat-Until-Success and `Option`\n",
    "\n",
    "We can now put both of the above together by preparing 5 qubits using `prepare_approx` and then attempting to distill a T state from them using `distill`, for some maximum number of attempts."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "bcc5acd4",
   "metadata": {},
   "outputs": [],
   "source": [
    "@guppy\n",
    "def t_state(attempts: int) -> Option[qubit]:\n",
    "    \"\"\"Create a T state using magic state distillation with a fixed number of attempts.\n",
    "\n",
    "    On success returns a qubit in a magic T state.\n",
    "\n",
    "    On failure (i.e. number of attempts are exceeded) returns nothing.\n",
    "    \"\"\"\n",
    "    if attempts > 0:\n",
    "        tgt = prepare_approx()\n",
    "        qs = array(prepare_approx() for _ in range(4))\n",
    "\n",
    "        q, success = distill(tgt, qs)\n",
    "        if success:\n",
    "            return some(q)\n",
    "        else:\n",
    "            # Discard the qubit and start over.\n",
    "            # Note, this could just as easily be a while loop!\n",
    "            discard(q)\n",
    "            return t_state(attempts - 1)\n",
    "\n",
    "    # We ran out of attempts.\n",
    "    return nothing()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d7945b35",
   "metadata": {},
   "source": [
    "Note the use of the `Option[qubit]` type for the result. Option types enable us to represent either a value, using `some(value)` or the absence of it, using `nothing()`. In this case it is a good way to represent failure or success without having to rely on errors. Before using the value if it exists, it needs to be unwrapped. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "9ff753c6",
   "metadata": {},
   "outputs": [],
   "source": [
    "attempts = 3\n",
    "\n",
    "\n",
    "@guppy\n",
    "def main() -> None:\n",
    "    option_t = t_state(comptime(attempts))\n",
    "    # Check whether the option contains a value.\n",
    "    if option_t.is_some():\n",
    "        # Unwrap the qubit.\n",
    "        t = option_t.unwrap()\n",
    "        state_result(\"t_state\", t)\n",
    "        discard(t)\n",
    "    else:\n",
    "        # Since qubits are linear, Option[qubit] is also linear, so we need to consume\n",
    "        # it either way.\n",
    "        option_t.unwrap_nothing()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "40a25f2d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n"
     ]
    }
   ],
   "source": [
    "shots = main.emulator(n_qubits=8).with_seed(1).run()\n",
    "\n",
    "# See the `Debugging with `state_result` statements` example notebook for more details\n",
    "# about state results.\n",
    "for states in shots.partial_state_dicts():\n",
    "    if \"t_state\" in states:\n",
    "        dist = states[\"t_state\"].state_distribution()\n",
    "        print(np.allclose(dist[0].state, np.array([0.2967+0.j, 0.8094-0.5068j]), rtol=1e-4))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0a8ee66c",
   "metadata": {},
   "source": [
    "A similar concept that is also available in the Guppy standard library is the `Either` type, which generalises the concept of having an optional value to having a value that could be of one of two types, either `left(value)` or `right(value)`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ce9f778c",
   "metadata": {},
   "source": [
    "## Selectively measuring arrays with `Option`\n",
    "\n",
    "Consider this attempt at measuring a subscript of an array of qubits:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "ee6a43c5",
   "metadata": {
    "tags": [
     "raises-exception"
    ]
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Error: Subscript consumed (at <In[7]>:3:12)\n",
      "  | \n",
      "1 | @guppy\n",
      "2 | def attempt_measure(qs: array[qubit, 10] @ owned) -> None:\n",
      "3 |     measure(qs[5])\n",
      "  |             ^^^^^ Cannot consume a subscript of `qs` with non-copyable type\n",
      "  |                   `array[qubit, 10]`\n",
      "\n",
      "Note: Subscripts on non-copyable types are only allowed to be borrowed, not\n",
      "consumed\n",
      "\n",
      "Guppy compilation failed due to 1 previous error\n"
     ]
    }
   ],
   "source": [
    "@guppy\n",
    "def attempt_measure(qs: array[qubit, 10] @ owned) -> None:\n",
    "    measure(qs[5])\n",
    "\n",
    "\n",
    "compiled = attempt_measure.compile()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d8861e5",
   "metadata": {},
   "source": [
    "As expected, this leads to an error because you can't consume subscripts of a linear array. \n",
    "\n",
    "However we can use arrays of type `array[Option[qubit], N]` to measure some qubits in an array without consuming the whole array at once by swapping qubits you want to measure with`nothing()`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "351ec80b",
   "metadata": {},
   "outputs": [],
   "source": [
    "n = guppy.nat_var(\"n\")\n",
    "\n",
    "\n",
    "@guppy\n",
    "def measure_mask(\n",
    "    qs: array[Option[qubit], n], mask: array[bool, n] @ owned\n",
    ") -> array[int, n]:\n",
    "    \"\"\"Measure all gubits in `qs` with a corresponding `True` index in `mask`.\"\"\"\n",
    "    res = array(-1 for _ in range(n))\n",
    "    idx = 0\n",
    "    for m in mask:\n",
    "        if m and qs[idx].is_some():\n",
    "            # `take` swaps an optional value with nothing().\n",
    "            q = qs[idx].take()\n",
    "            res[idx] = int(measure(q.unwrap()))\n",
    "        idx += 1\n",
    "    return res\n",
    "\n",
    "\n",
    "@guppy\n",
    "def main() -> None:\n",
    "    qs = array(some(qubit()) for _ in range(3))\n",
    "    mask = array(True, False, True)\n",
    "    measure_mask(qs, mask)\n",
    "    # We need to consume the array of options at some point, as it can't be leaked.\n",
    "    for q_opt in qs:\n",
    "        if q_opt.is_some():\n",
    "            q = q_opt.unwrap()\n",
    "            discard(q)\n",
    "        else:\n",
    "            q_opt.unwrap_nothing()\n",
    "\n",
    "\n",
    "main.compile();"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv",
   "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.11.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}