{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c0439942",
   "metadata": {},
   "source": [
    "# Lesson 2b: Writing Workers (Guppy)\n",
    "\n",
    "```{important}\n",
    "This is an alternative lesson to [Lesson 2a](pytket_graph.ipynb).\n",
    "It uses [Guppy](https://guppylang.org) to construct the quantum computation.\n",
    "```\n",
    "\n",
    "In Lesson 2a you wrote a worker that constructs a symbolic quantum circuit and running it on a simulator.\n",
    "In this example we're going to extend your \"my_example_worker\" with Guppy functions.\n",
    "1. Define a simple parametric Guppy program\n",
    "2. Provide parameters at runtime\n",
    "\n",
    "```{note}\n",
    "Guppy is a programming language for Quantum Computers.\n",
    "All code in a Guppy program will be executed in coherence time on the controller.\n",
    "Since it is embedded in Python, we can use the same worker mechanism as before.\n",
    "\n",
    "```\n",
    "[Lesson 3](./storage_and_executors.ipynb) will proceed with the pytket functionality, so you may skip ahead.\n",
    "\n",
    "## Prerequisite\n",
    "\n",
    "We're going to extend the existing worker so we assume you have run the prerequisites from before.\n",
    "\n",
    "As a reminder this was (see [2a](pytket_graph.ipynb#Prerequisite)):\n",
    "\n",
    "```bash\n",
    "uv init\n",
    "uv add tierkreis pytket pytket-qiskit sympy ruff\n",
    "uv run tkr project init\n",
    "uv run tkr init worker -n my_example_worker\n",
    "```\n",
    "We're adding the dependencies for Guppy:\n",
    "```bash\n",
    "uv add guppylang hugr\n",
    "```\n",
    "\n",
    "## Defining the tasks\n",
    "\n",
    "As before we assume your implementing `impl.py` of your previously defined worker"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16ecff0a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# The worker is added here for validity of the example\n",
    "from tierkreis import Worker\n",
    "\n",
    "worker = Worker(\"my_example_worker\")\n",
    "...  # Your previous tasks will be here"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1416b27e",
   "metadata": {},
   "source": [
    "In Guppy, you write programs by using the `@guppy` annotation and then calling compile.\n",
    "You can do the same from inside a worker function, here we will dynamically construct a n-qubit GHZ state."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19da7c1c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from guppylang import comptime, guppy\n",
    "from guppylang.std.builtins import array, result\n",
    "from guppylang.std.quantum import cx, h, measure_array, qubit\n",
    "from hugr.package import Package\n",
    "\n",
    "\n",
    "@worker.task()\n",
    "def ghz(size: int) -> Package:\n",
    "    n = guppy.nat_var(\"n\")\n",
    "\n",
    "    @guppy\n",
    "    def build_ghz_state(q: array[qubit, n]) -> None:  # type: ignore\n",
    "        h(q[0])\n",
    "        for i in range(n - 1):  # type: ignore\n",
    "            cx(q[i], q[i + 1])\n",
    "\n",
    "    @guppy\n",
    "    def main() -> None:\n",
    "        q = array(qubit() for _ in range(comptime(size)))  # type: ignore\n",
    "        build_ghz_state(q)\n",
    "\n",
    "        result(\"c\", measure_array(q))\n",
    "\n",
    "    return main.compile()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ad3bf915",
   "metadata": {},
   "source": [
    "As before you can provide inputs to the function.\n",
    "```{important}\n",
    "You must call `compile()` on the main function you wish to run as a task.\n",
    "This will generate a serilaized version of the program (a Hugr), which can be used as value in Tierkreis. The corresponding type is `hugr.Package`.\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4bb179a7",
   "metadata": {},
   "source": [
    "\n",
    "## Generating stubs\n",
    "\n",
    "As before generate the api from the cli\n",
    "```bash\n",
    "uv run tkr init stubs\n",
    "```\n",
    "\n",
    "## Using the tasks\n",
    "\n",
    "Now you can use the newly declared tasks in a graph similar to how you used the `builtin` functionality or other tasks.\n",
    "You have to import the task API from the worker first which you then can use with a task node.\n",
    "First we declare the graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "acf47c53",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Constructing, put into graphs/main.py\n",
    "from tierkreis.builder import Graph\n",
    "from tierkreis.controller.data.models import TKR\n",
    "\n",
    "\n",
    "graph = Graph(TKR[int], TKR[Package])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c16d4459",
   "metadata": {},
   "source": [
    "and then add the tasks:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5e445a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Constructing, put into graphs/main.py\n",
    "from my_example_worker import ghz  # noqa: F811\n",
    "\n",
    "program = graph.task(ghz(graph.inputs))\n",
    "workflow = graph.finish_with_outputs(program)  # type: ignore"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6d4f8e3f",
   "metadata": {},
   "source": [
    "## Running the graph\n",
    "As before you know can run the graph, the circuit we have defined already above. The result is an binary represantion of the program."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "06862811",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Running, put into graphs/main.py\n",
    "from uuid import UUID\n",
    "\n",
    "\n",
    "from tierkreis.controller import run_graph\n",
    "from tierkreis.executor import ShellExecutor\n",
    "from tierkreis.storage import FileStorage, read_outputs\n",
    "\n",
    "\n",
    "storage = FileStorage(workflow_id=UUID(int=12347), name=\"Guppy example graph\")\n",
    "storage.clean_graph_files()\n",
    "executor = ShellExecutor(registry_path=None, workflow_dir=storage.workflow_dir)\n",
    "run_graph(storage, executor, workflow, 3)  # single import\n",
    "output = read_outputs(workflow, storage)\n",
    "print(output)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "tierkreis",
   "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.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}