{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "39011d5f",
   "metadata": {},
   "source": [
    "# Complex types, optional values\n",
    "So far we have only build graphs with a single input and output and defined everything else as constants.\n",
    "Tierkreis allows for multiple inputs and outputs similar to python, by defining tuples of multiple fields.\n",
    "As an invariance, a graph will still always have exactly one output node which also mimics python.\n",
    "\n",
    "In this example we will look at defining graph which compiles and runs a single circuit on an IBM simulator.\n",
    "For this you will need valid IBMQ credentials.\n",
    "\n",
    "## Inputs and Outputs\n",
    "\n",
    "In tierkreis, inputs and outputs are defined as named tuples of `TKR` types.\n",
    "Were building a graph circuit -> int  -> (BackendResult, int)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "def19061",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install tierkreis pytket"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "86970dbb",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import NamedTuple\n",
    "\n",
    "from tierkreis.controller.data.models import TKR, OpaqueType\n",
    "\n",
    "\n",
    "class IBMInput(NamedTuple):\n",
    "    circuit: TKR[OpaqueType[\"pytket._tket.circuit.Circuit\"]]  # noqa: F821\n",
    "    n_shots: TKR[int]\n",
    "\n",
    "\n",
    "class IBMOutput(NamedTuple):\n",
    "    results: TKR[OpaqueType[\"pytket._tket.backends.BackendResult\"]]  # noqa: F821\n",
    "    n_shots: TKR[int]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "88dbc3f3",
   "metadata": {},
   "source": [
    "Now we construct the graph using the `tkr-aer-worker`.\n",
    "Note how we address inputs by their name `g.inputs.circuit` and for the output we construct an output object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5c075ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "from tierkreis.aer_worker import get_compiled_circuit, submit_single\n",
    "from tierkreis.builder import GraphBuilder\n",
    "\n",
    "\n",
    "def compile_run_single() -> GraphBuilder[IBMInput, IBMOutput]:\n",
    "    g = GraphBuilder(IBMInput, IBMOutput)\n",
    "\n",
    "    compiled_circuit = g.task(\n",
    "        get_compiled_circuit(\n",
    "            circuit=g.inputs.circuit,\n",
    "            optimisation_level=g.const(2),\n",
    "        ),\n",
    "    )\n",
    "    res = g.task(submit_single(compiled_circuit, g.inputs.n_shots))\n",
    "    g.outputs(IBMOutput(res, g.inputs.n_shots))\n",
    "    return g"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a23932fe",
   "metadata": {},
   "source": [
    "We now can run this code as before. Inputs currently still have to be defined as inputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "031538db",
   "metadata": {},
   "outputs": [],
   "source": [
    "from uuid import UUID\n",
    "\n",
    "from pytket._tket.circuit import Circuit\n",
    "\n",
    "from tierkreis.consts import PACKAGE_PATH\n",
    "from tierkreis.controller import run_graph\n",
    "from tierkreis.executor import UvExecutor\n",
    "from tierkreis.storage import FileStorage, read_outputs\n",
    "\n",
    "storage = FileStorage(UUID(int=109), do_cleanup=True, name=\"ibmq_example\")\n",
    "executor = UvExecutor(PACKAGE_PATH.parent / \"tierkreis_workers\", storage.logs_path)\n",
    "\n",
    "\n",
    "def ghz() -> Circuit:\n",
    "    circ1 = Circuit(2)\n",
    "    circ1.H(0)\n",
    "    circ1.CX(0, 1)\n",
    "    circ1.measure_all()\n",
    "    return circ1\n",
    "\n",
    "\n",
    "inputs = {\n",
    "    \"circuit\": ghz(),\n",
    "    \"n_shots\": 1024,\n",
    "    \"backend\": \"ibm_torino\",\n",
    "}\n",
    "graph = compile_run_single()\n",
    "run_graph(\n",
    "    storage,\n",
    "    executor,\n",
    "    graph,\n",
    "    inputs,\n",
    ")\n",
    "res = read_outputs(graph, storage)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d7baa21a",
   "metadata": {},
   "source": [
    "## Optional Values\n",
    "\n",
    "Tierkreis can also deal with optional values which are indicated python style bny providing the type hint `type | None = None`.\n",
    "Previously we have defined the optimization level by setting a constant.\n",
    "Inspecting the type declaration of `compile_circuit_ibmq` reveals the following type: `optimisation_level: TKR[str] | None = None`\n",
    "which means we can set it optionally.\n",
    "We will also pull this up into the inputs:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5b38b249",
   "metadata": {},
   "outputs": [],
   "source": [
    "class IBMInputOptional(NamedTuple):\n",
    "    circuit: TKR[OpaqueType[\"pytket._tket.circuit.Circuit\"]]  # noqa: F821\n",
    "    n_shots: TKR[int]\n",
    "    optimisation_level: TKR[int] | None = None  # Optional input\n",
    "\n",
    "\n",
    "def compile_optional() -> GraphBuilder[IBMInputOptional, IBMOutput]:\n",
    "    g = GraphBuilder(IBMInputOptional, IBMOutput)\n",
    "\n",
    "    compiled_circuit = g.task(\n",
    "        get_compiled_circuit(\n",
    "            circuit=g.inputs.circuit,\n",
    "            optimisation_level=g.inputs.optimisation_level,\n",
    "        ),\n",
    "    )\n",
    "    res = g.task(submit_single(compiled_circuit, g.inputs.n_shots))\n",
    "    g.outputs(IBMOutput(res, g.inputs.n_shots))  # type: ignore\n",
    "    return g"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f589bec6",
   "metadata": {},
   "source": [
    "You can now run the the graph with or without th optional input.\n",
    "Note, that not providing all inputs will produce some error messages."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a941c26f",
   "metadata": {},
   "outputs": [],
   "source": [
    "inputs = {\n",
    "    \"circuit\": ghz(),\n",
    "    \"n_shots\": 1024,\n",
    "    # \"optimisation_level\": 1, Uncomment to provide the optional input\n",
    "}\n",
    "\n",
    "graph = compile_optional()\n",
    "storage.clean_graph_files()\n",
    "run_graph(\n",
    "    storage,\n",
    "    executor,\n",
    "    graph,\n",
    "    inputs,\n",
    ")\n",
    "res = read_outputs(graph, storage)"
   ]
  }
 ],
 "metadata": {
  "execution": {
   "timeout": 120
  },
  "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.13.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}