{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c0439942",
   "metadata": {},
   "source": [
    "# Lesson 2: Writing Workers\n",
    "\n",
    "In the previous example you wrote a worker only using Tierkreis base functionality.\n",
    "One of the main benefits of using Tierkreis is that you can easily transform existing code and gain the benefits of Tierkeis, e.g., checkpointing and repeteability.\n",
    "\n",
    "In this example we're going to look at a common problem and use as tasks in a graph.\n",
    "For this we're writing a custom worker \"my_example_worker\".\n",
    "We're going to use [pytket](https://docs.quantinuum.com/tket/api-docs/) to:\n",
    "1. Define a symbolic circuit\n",
    "2. Substitute its symbolic parameters at runtime\n",
    "\n",
    "Compiling and running circuits will be part of [Lesson 3](./storage_and_executors.ipynb).\n",
    "\n",
    "## Prerequisite\n",
    "\n",
    "If you haven't done so, set up a Tierkreis project and install the additional dependencies for pytket and sympy.\n",
    "Ruff is an optional dependency for formatting the autogenerated code.\n",
    "\n",
    "```bash\n",
    "uv init\n",
    "uv add tierkreis pytket pytket-qiskit sympy ruff\n",
    "uv run tkr project init\n",
    "```\n",
    "\n",
    "## Setting up the worker\n",
    "\n",
    "Using the the Tierkreis cli is the easiest the way to set up a worker:\n",
    "\n",
    "```bash\n",
    "uv run tkr init worker -n my_example_worker\n",
    "```\n",
    "\n",
    "Which will set up the packages in a convenient way for you.\n",
    "```{info}\n",
    "If you don't want to use the cli, you need a way to provide the worker API to your graph construction.\n",
    "You can find more information [here](../worker/index.md)\n",
    "```\n",
    "\n",
    "The cli creates the folder structure as seen in the first lesson.\n",
    "There are two files in `tkr/workers/my_example_worker` which you need to care about which are:\n",
    "- `tkr_my_example_worker_impl/impl.py` here you will implement your workers functionality\n",
    "- `api/api.py` is an autogenerated file which contains the task definitions we will use in the graph.\n",
    "\n",
    "In `impl.py` you will find the following:\n",
    "\n",
    "```python\n",
    "worker = Worker(\"my_example_worker\")\n",
    "\n",
    "@worker.task()\n",
    "def your_worker_task(value: int) -> int:\n",
    "    return value\n",
    "```\n",
    "\n",
    "As you can see, generating a task for Tierkeis simply means adding `@worker.task()` to a function.\n",
    "We're going to use the same pattern to expose `pytket`s functionality to in our new worker.\n",
    "\n",
    "## Defining the tasks\n",
    "\n",
    "```{important}\n",
    "Not all cells in the following section are not necessary to run this code example.\n",
    "These are example tasks that are meant to sit in your own copy of the `my_example_worker`s `impl.py`.\n",
    "If you're running this example from the repository, it will contain a copy of the worker already installed.\n",
    "```\n",
    "\n",
    "For the following assume we're going to use a simple quantum circuits with three symbols $a,b,c$.\n",
    "Writing  this as a worker task means wrapping it with a task function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16ecff0a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# The worker is added here for validity of the example\n",
    "# Put the task into impl.py\n",
    "from pytket.circuit import Circuit, fresh_symbol\n",
    "from tierkreis import Worker\n",
    "\n",
    "worker = Worker(\"pytket_example_worker\")\n",
    "\n",
    "\n",
    "@worker.task()\n",
    "def symbolic_circuit() -> Circuit:\n",
    "    a = fresh_symbol(\"a\")\n",
    "    b = fresh_symbol(\"b\")\n",
    "    c = fresh_symbol(\"c\")\n",
    "    circ = Circuit(3)\n",
    "    circ.Rz(a, 0)\n",
    "    circ.Rz(b, 0)\n",
    "    circ.Rz(c, 0)\n",
    "    circ.measure_all()\n",
    "    return circ"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1416b27e",
   "metadata": {},
   "source": [
    "Using plain `pytket` you could know substitute the symbols like so:\n",
    "```python\n",
    "# For explanation only, this will be a task\n",
    "from sympy import Symbol\n",
    "\n",
    "circ.symbol_substitution({Symbol(\"a\"): -1, Symbol(\"b\"): 0, Symbol(\"c\"): 1})\n",
    "```\n",
    "writing this as a worker task means wrapping it with a task function that has a circuit input and output.\n",
    "```{important}\n",
    "Use type hints so that Tierkreis can validate the task during construction.\n",
    "Since values are represented by edges, we need a return value, its not sufficient to mutate the state.\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19da7c1c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Put the task into impl.py\n",
    "from sympy import Symbol\n",
    "\n",
    "\n",
    "@worker.task()\n",
    "def substitute(circuit: Circuit, a: float, b: float, c: float) -> Circuit:\n",
    "    circuit.symbol_substitution({Symbol(\"a\"): a, Symbol(\"b\"): b, Symbol(\"c\"): c})\n",
    "    return circuit"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ad3bf915",
   "metadata": {},
   "source": [
    "Similarly we could now define other tasks using `pytket` and `@worker.task()`\n",
    "e.g.,compilation and simulation."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4bb179a7",
   "metadata": {},
   "source": [
    "\n",
    "## Generating stubs\n",
    "\n",
    "To generate the APIs for all workers, you can use the cli with\n",
    "```bash\n",
    "uv run tkr init stubs\n",
    "```\n",
    "\n",
    "Depending on your development environment it might be necessary to resatart your language server or `uv sync --all-extras` to pick up the update changes.\n",
    "\n",
    "\n",
    "### Opaque Types\n",
    "Tierkreis can use any type that is serializable as in and outputs, e.g., the `Circuit` type from `pytket` library.\n",
    "To make such types available without bleeding dependencies into graph code, Tiekreis wraps them as `OpaqueType` with a reference to the original implementation.\n",
    "In this example the `circuit` inputs of the tasks would be\n",
    "\n",
    "```python\n",
    "circuit: TKR[OpaqueType[\"pytket._tket.circuit.Circuit\"]] \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.\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 typing import NamedTuple\n",
    "from tierkreis.builder import Graph\n",
    "from tierkreis.controller.data.models import TKR\n",
    "\n",
    "\n",
    "class PytketInputs(NamedTuple):\n",
    "    a: TKR[float]\n",
    "    b: TKR[float]\n",
    "    c: TKR[float]\n",
    "\n",
    "\n",
    "graph = Graph(PytketInputs, TKR[Circuit])"
   ]
  },
  {
   "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 substitute, symbolic_circuit  # noqa: F811\n",
    "\n",
    "circuit = graph.task(symbolic_circuit())\n",
    "substituted = graph.task(\n",
    "    substitute(circuit, graph.inputs.a, graph.inputs.a, graph.inputs.a)  # type: ignore\n",
    ")\n",
    "workflow = graph.finish_with_outputs(substituted)  # 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"
   ]
  },
  {
   "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=12346), name=\"Pytket example graph\")\n",
    "storage.clean_graph_files()\n",
    "executor = ShellExecutor(registry_path=None, workflow_dir=storage.workflow_dir)\n",
    "run_graph(storage, executor, workflow, {\"a\": -1, \"b\": 0, \"c\": 1})\n",
    "result = read_outputs(workflow, storage)\n",
    "print(result)"
   ]
  }
 ],
 "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
}