{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "f94615ed",
   "metadata": {},
   "source": [
    "# Experimenting with Parallel Teleport on Neutral Atom Device\n",
    "\n",
    "This notebook shows how you can use the `NeutralAtomDevice` class to visualize and simulate a scaled, parallel teleport sample program.\n",
    "\n",
    "First, set up the environment."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae94a027",
   "metadata": {},
   "outputs": [],
   "source": [
    "from qdk import init, TargetProfile, code, Result\n",
    "from qdk.qsharp import compile, circuit\n",
    "from qdk.simulation import NeutralAtomDevice, NoiseConfig\n",
    "from qdk.widgets import Circuit, Histogram\n",
    "\n",
    "init(target_profile=TargetProfile.Base)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "df5d6387",
   "metadata": {},
   "source": [
    "Next, we use Q# to define a `ParalellTeleport` that uses distinct sets of qubits to perform a measurement deferred teleportation. Using different qubits rather than the same qubits should allow later device-specific compilation steps to identify parallelism and run those teleports across the qubits at the same time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0d435592",
   "metadata": {
    "vscode": {
     "languageId": "qsharp"
    }
   },
   "outputs": [],
   "source": [
    "%%qsharp\n",
    "\n",
    "operation ParallelTeleport(instances : Int) : Result[] {\n",
    "    // Partitions the teleport instances across rows and columns on the machine\n",
    "    let cols = if instances >= 12 { 36 } else { instances * 3 };\n",
    "    let rows = (instances + 11) / 12;  // 1 to 12 = 1, 13 to 24 = 2, etc.\n",
    "\n",
    "    use qubits = Qubit[instances * 3];\n",
    "    mutable results : Result[] = [];\n",
    "\n",
    "    for i in 0..instances-1 {\n",
    "        let rowId = i / 12;\n",
    "        let colId = (i % 12) * 3;\n",
    "        let idx = colId + (rowId * 36);\n",
    "\n",
    "        // Prep state on msg qubit\n",
    "        if i % 4 == 1 {\n",
    "            X(qubits[idx]);\n",
    "        } elif i % 4 == 2 {\n",
    "            H(qubits[idx]);\n",
    "        } elif i % 4 == 3 {\n",
    "            SX(qubits[idx]);\n",
    "        }\n",
    "\n",
    "        Teleport(qubits[idx], qubits[idx + 1], qubits[idx + 2]);\n",
    "\n",
    "        // Reverse state prep on Bob's qubit\n",
    "        if i % 4 == 1 {\n",
    "            X(qubits[idx + 2]);\n",
    "        } elif i % 4 == 2 {\n",
    "            H(qubits[idx + 2]);\n",
    "        } elif i % 4 == 3 {\n",
    "            X(qubits[idx + 2]);\n",
    "            SX(qubits[idx + 2]);\n",
    "        }\n",
    "        results += [MResetZ(qubits[idx + 2])];\n",
    "        ResetAll([qubits[idx], qubits[idx + 1]]);\n",
    "    }\n",
    "\n",
    "    return results;\n",
    "}\n",
    "\n",
    "operation Teleport(msg : Qubit, alice : Qubit, bob : Qubit) : Unit {\n",
    "    // Create some entanglement that we can use to send our message.\n",
    "    H(alice);\n",
    "    CNOT(alice, bob);\n",
    "\n",
    "    // Encode the message into the entangled pair.\n",
    "    CNOT(msg, alice);\n",
    "    H(msg);\n",
    "\n",
    "    CNOT(alice, bob);\n",
    "    Controlled Z([msg], bob);\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d86e19f4",
   "metadata": {},
   "source": [
    "You can verify the circuit for two parallel teleports using the `Circuit` widget."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b815acd4",
   "metadata": {},
   "outputs": [],
   "source": [
    "Circuit(circuit(code.ParallelTeleport, 2))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c342d7e9",
   "metadata": {},
   "source": [
    "Now, let's visualize 120 instances of teleportation in parallel on the `NeutralAtomDevice`. This will decompose the program to the device's native gate set and schedule the operations with movement, showing how it might execute on such a device."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2422917",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compile the code into a high-level program\n",
    "program = compile(code.ParallelTeleport, 120)\n",
    "\n",
    "# Visualize a trace of running the program on the target machine,\n",
    "# decomposing into native gates and scheduling parallel operations\n",
    "device = NeutralAtomDevice()\n",
    "device.show_trace(program)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "87b8e819",
   "metadata": {},
   "source": [
    "Because the teleport algorithm uses only gates from the Clifford group, we can efficiently simulate a large number of shots of the same program with the Clifford simulator. We'll use a small helper function to interpret the results, which we expect to be all zeros. Without noise, \"Correct\" should be the only bar in the resulting Histogram."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "698849c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "results = device.simulate(program, shots=1000, type=\"clifford\")\n",
    "Histogram([\"Loss\" if any(r == Result.Loss for r in shot) else \"Flip\" if any(r == Result.One for r in shot) else \"Correct\" for shot in results])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ce124bf0",
   "metadata": {},
   "source": [
    "You can also perform this simulation with noise, shown here by configuring a 0.1% chance of loss on the SX gate. Now we'll see a mix of correct and incorrect results in the output of the simulation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "774df3a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Now with noise.\n",
    "noise = NoiseConfig()\n",
    "noise.sx.loss = 0.001\n",
    "results = device.simulate(program, shots=1000, noise=noise, type=\"clifford\")\n",
    "Histogram([\"Loss\" if any(r == Result.Loss for r in shot) else \"Flip\" if any(r == Result.One for r in shot) else \"Correct\" for shot in results])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 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.11.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
