{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "8b376019",
   "metadata": {},
   "source": [
    "# Debugging with `state_result` statements\n",
    "\n",
    "**Download this notebook - {nb-download}`state_results.ipynb`**\n",
    "\n",
    "In this example we will demonstrate the use of `state_result` statements for surfacing the simulator state while running a Guppy program with Selene. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "aa81f6b0",
   "metadata": {},
   "outputs": [],
   "source": [
    "from guppylang import guppy\n",
    "from guppylang.std.builtins import array\n",
    "from guppylang.std.debug import state_result\n",
    "from guppylang.std.quantum import qubit, discard, discard_array, cx, h, measure, x\n",
    "\n",
    "from hugr.qsystem.result import QsysResult\n",
    "\n",
    "from selene_sim import build, Quest\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "np.set_printoptions(precision=4, suppress=True, linewidth=120)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aed12879",
   "metadata": {},
   "source": [
    "Consider the following Guppy program where we create a 3-qubit GHZ state in two different ways that should be unitarily equivalent. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "dba2658d",
   "metadata": {},
   "outputs": [],
   "source": [
    "@guppy\n",
    "def ghz_state(q: array[qubit, 3]) -> None:\n",
    "    h(q[0])\n",
    "    cx(q[0], q[1])\n",
    "    # Apply CX to second and third qubit.\n",
    "    cx(q[1], q[2])\n",
    "\n",
    "\n",
    "@guppy\n",
    "def ghz_state_alternative(q: array[qubit, 3]) -> None:\n",
    "    h(q[0])\n",
    "    cx(q[0], q[1])\n",
    "    # Apply CX to first and third qubit.\n",
    "    cx(q[0], q[2])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2bc04c3f",
   "metadata": {},
   "source": [
    "In order to check whether the resulting state vectors are equal, we can use the `state_result` function. As inputs it takes a tag (just as `result` statements do) and then either one array of qubits or a variable number of single qubit arguments. The order of arguments later determines the order of state vector outputs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "23d0ac6e",
   "metadata": {},
   "outputs": [],
   "source": [
    "@guppy\n",
    "def main() -> None:\n",
    "    qs = array(qubit() for _ in range(3))\n",
    "    ghz_state(qs)\n",
    "    state_result(\"ghz\", qs)\n",
    "    discard_array(qs)\n",
    "\n",
    "    qs = array(qubit() for _ in range(3))\n",
    "    ghz_state_alternative(qs)\n",
    "    state_result(\"ghz_alternative\", qs)\n",
    "    discard_array(qs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "79ecfb4d",
   "metadata": {},
   "source": [
    "Let's now extract the requested state from the simulator. This can be done by calling `extract_states_dict` on the simulator plugin for each shot to get a dictionary mapping tags to states, and then retrieving the results from the state using `get_single_state`. \n",
    "\n",
    "Here our `ghz_state` and `ghz_state_alternative` functions deterministically prepare a single quantum state.\n",
    "\n",
    "Note that this is currently only possible using the `Quest` simulator plugin with the default `SimpleRuntime` and `IdealErrorModel`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "1d27f005",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "State `ghz1`: [0.7071+0.j 0.    +0.j 0.    +0.j 0.    +0.j 0.    +0.j 0.    +0.j 0.    +0.j 0.7071+0.j]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "State `ghz2`: [0.7071+0.j 0.    +0.j 0.    +0.j 0.    +0.j 0.    +0.j 0.    +0.j 0.    +0.j 0.7071+0.j]\n"
     ]
    }
   ],
   "source": [
    "shots = main.emulator(n_qubits=3).statevector_sim().with_seed(42).run()\n",
    "\n",
    "for states in shots.partial_state_dicts():\n",
    "    state_vector1 = states[\"ghz\"].as_single_state()\n",
    "    print(f\"State `ghz1`: {state_vector1}\")\n",
    "\n",
    "    state_vector2 = states[\"ghz_alternative\"].as_single_state()\n",
    "    print(f\"State `ghz2`: {state_vector2}\")\n",
    "\n",
    "    assert np.allclose(state_vector1, state_vector2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "893c11a0",
   "metadata": {},
   "source": [
    "We can see that the state vector is the same for both GHZ functions!\n",
    "\n",
    "It is important to bear in mind that the state result is simply the internal state of the simulator plugin, so the output at any given point in the middle of the program could be very different on hardware, where various optimisations and different scheduling likely apply.\n",
    "\n",
    "If you chose to request a subset of the qubits making up the overall state, you should also be aware that qubits in this subset could be entangled with other qubits in the system, so the qubits not requested in the state result are traced out to provide a probabilistic set of state vectors. For example, if we only request the state of any two out of the three qubits in the GHZ state, we expect to see either $\\ket{00}$ or $\\ket{11}$ with a probability of $\\frac{1}{2}$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "063ea31e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[1.+0.j 0.+0.j 0.+0.j 0.+0.j] -> 0.5, [0.+0.j 0.+0.j 0.+0.j 1.+0.j] -> 0.5\n",
      "[1.+0.j 0.+0.j 0.+0.j 0.+0.j] -> 0.5, [0.+0.j 0.+0.j 0.+0.j 1.+0.j] -> 0.5\n",
      "[1.+0.j 0.+0.j 0.+0.j 0.+0.j] -> 0.5, [0.+0.j 0.+0.j 0.+0.j 1.+0.j] -> 0.5\n"
     ]
    }
   ],
   "source": [
    "@guppy\n",
    "def main() -> None:\n",
    "    qs = array(qubit() for _ in range(3))\n",
    "    ghz_state(qs)\n",
    "    state_result(\"subset01\", qs[0], qs[1])\n",
    "    state_result(\"subset02\", qs[0], qs[2])\n",
    "    state_result(\"subset12\", qs[1], qs[2])\n",
    "    discard_array(qs)\n",
    "\n",
    "\n",
    "runner = build(main.compile())\n",
    "\n",
    "shots = main.emulator(n_qubits=4).statevector_sim().with_seed(42).run()\n",
    "\n",
    "\n",
    "for states in shots.partial_state_dicts():\n",
    "    state01 = states[\"subset01\"].state_distribution()\n",
    "    state02 = states[\"subset02\"].state_distribution()\n",
    "    state12 = states[\"subset12\"].state_distribution()\n",
    "\n",
    "    print(\n",
    "        f\"{state01[0].state} -> {round(state01[0].probability, 2)}, {state01[1].state} -> {round(state01[1].probability, 2)}\"\n",
    "    )\n",
    "    print(\n",
    "        f\"{state02[0].state} -> {round(state02[0].probability, 2)}, {state02[1].state} -> {round(state02[1].probability, 2)}\"\n",
    "    )\n",
    "    print(\n",
    "        f\"{state12[0].state} -> {round(state12[0].probability, 2)}, {state12[1].state} -> {round(state12[1].probability, 2)}\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "82e40c32",
   "metadata": {},
   "source": [
    "Also keep in mind that state results can vary between shots depending on control flow."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "2ce50362",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{0: 58, 1: 42}"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "@guppy\n",
    "def main() -> None:\n",
    "    q = qubit()\n",
    "    h(q)\n",
    "    q1 = qubit()\n",
    "    q2 = qubit()\n",
    "    x(q2)\n",
    "    if measure(q):\n",
    "        state_result(\"state\", q2)\n",
    "    else:\n",
    "        state_result(\"state\", q1)\n",
    "    discard(q1)\n",
    "    discard(q2)\n",
    "\n",
    "shots = main.emulator(n_qubits=4).statevector_sim().with_seed(2).with_shots(100).run()\n",
    "\n",
    "\n",
    "state_counts = {0: 0, 1: 0}\n",
    "for state_dict in shots.partial_state_dicts():\n",
    "    state = state_dict[\"state\"].state_distribution()\n",
    "    s = state[0].state\n",
    "    if np.allclose(s, [0, 1]):\n",
    "        state_counts[1] += 1\n",
    "    elif np.allclose(s, [1, 0]):\n",
    "        state_counts[0] += 1\n",
    "\n",
    "print(state_counts)"
   ]
  }
 ],
 "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
}