{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "5027ce86",
   "metadata": {},
   "source": [
    "# Step, debug, and replay\n",
    "\n",
    "`pipe.step(dt=...)` advances a Retriever graph **one tick at a time, in your own\n",
    "Python process** — no backend, no separate processes, so a normal debugger stops\n",
    "right inside `Flow.step()`. This notebook steps a small graph, reads back a\n",
    "`StepResult`, inspects the validated IR with `pipe.validate()`, and records a run\n",
    "to disk so it can be replayed deterministically. Every cell runs in-process with\n",
    "only `retriever-core` — no camera, GPU, or robot."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "916b2888",
   "metadata": {},
   "source": [
    "> **Running in Colab?** The next cell installs `retriever-core`. From a source\n",
    "> checkout (or once it's already installed) the install is skipped."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "38b9b2eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Colab setup: install retriever-core only if it isn't importable yet.\n",
    "try:\n",
    "    import retriever  # noqa: F401\n",
    "except ImportError:  # pragma: no cover\n",
    "    import subprocess\n",
    "    import sys\n",
    "\n",
    "    subprocess.run(\n",
    "        [sys.executable, \"-m\", \"pip\", \"install\", \"retriever-core\"], check=True\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "624e0fee",
   "metadata": {},
   "source": [
    "## A graph to debug\n",
    "\n",
    "Three tiny Flows: a `Counter` source on a clock, a `Doubler` that fires on each\n",
    "arrival, and a `Sink` that keeps the last value it saw. This is an ordinary\n",
    "Retriever graph — the only difference is that we'll drive it with `step()`\n",
    "instead of handing it to a backend."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e7f6a9cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "from retriever.flow import Flow, Latest, Pipeline, Rate, Trigger, io\n",
    "\n",
    "\n",
    "@io\n",
    "class Value:\n",
    "    value: int\n",
    "\n",
    "\n",
    "class Counter(Flow[None, Value]):\n",
    "    def reset(self) -> None:\n",
    "        self.count = 0\n",
    "\n",
    "    def step(self, _) -> Value:\n",
    "        self.count += 1\n",
    "        return Value(value=self.count)\n",
    "\n",
    "\n",
    "class Doubler(Flow[Value, Value]):\n",
    "    def step(self, input: Value) -> Value:\n",
    "        return Value(value=input.value * 2)  # set a breakpoint here\n",
    "\n",
    "\n",
    "class Sink(Flow[Value, None]):\n",
    "    def reset(self) -> None:\n",
    "        self.last = None\n",
    "\n",
    "    def step(self, input: Value) -> None:\n",
    "        self.last = input.value\n",
    "        return None\n",
    "\n",
    "\n",
    "pipe = Pipeline(\"step_debug_replay\")\n",
    "with pipe:\n",
    "    counter = Counter() @ Rate(hz=10)\n",
    "    doubler = Doubler() @ Trigger(\"value\")\n",
    "    sink = Sink() @ Rate(hz=10)\n",
    "    pipe.connect(counter, doubler, sync=Latest())\n",
    "    pipe.connect(doubler, sink, sync=Latest())\n",
    "\n",
    "print(\"built:\", pipe.get_name(), \"with\", len(pipe.get_handles()), \"nodes\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d5ea18d5",
   "metadata": {},
   "source": [
    "## `pipe.step(dt=...)` advances one tick\n",
    "\n",
    "`step(dt=0.1)` runs one `sample -> step -> publish` pass over the whole graph and\n",
    "advances a logical clock by `dt`. It returns synchronously, so a breakpoint in\n",
    "any `Flow.step()` stops here in this process — no remote debugger to attach."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6dcd36c6",
   "metadata": {},
   "outputs": [],
   "source": [
    "pipe.reset()  # gym-style: clear buffers and call Flow.reset() on every node\n",
    "\n",
    "first_now = None\n",
    "for i in range(4):\n",
    "    result = pipe.step(dt=0.1)\n",
    "    if first_now is None:\n",
    "        first_now = result.now\n",
    "    elapsed = round(result.now - first_now, 2)\n",
    "    print(f\"tick {i}: t=+{elapsed:.1f}s  executed={sorted(result.executed)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c8c2da72",
   "metadata": {},
   "source": [
    "Every tick advances logical time by `dt` and, because both clocks fire each tick,\n",
    "all three nodes execute. The debugger treats `step()` like any other function."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5fb8c635",
   "metadata": {},
   "source": [
    "## Read back a `StepResult`\n",
    "\n",
    "Each `step()` returns a `StepResult` — a frozen record of what that tick did:\n",
    "`.executed` (node ids that ran), `.outputs` (each node's returned object), and\n",
    "`.inputs` (what each node consumed). No printing inside a Flow required; you read\n",
    "the tick's effects straight off the result."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e71fc866",
   "metadata": {},
   "outputs": [],
   "source": [
    "pipe.reset()\n",
    "result = pipe.step(dt=0.1)\n",
    "\n",
    "print(\"type    :\", type(result).__name__)\n",
    "print(\"executed:\", sorted(result.executed))\n",
    "for node_id in sorted(result.outputs):\n",
    "    print(f\"output[{node_id}] = {result.outputs[node_id]}\")\n",
    "print(\"sink.last =\", sink.flow.last)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "602d74c0",
   "metadata": {},
   "source": [
    "`Counter` emitted `value=1`; `Doubler` turned it into `value=2`; `Sink` returned\n",
    "`None` but stashed the value in its own state — exactly what a breakpoint would\n",
    "have shown you mid-tick."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "30d7c427",
   "metadata": {},
   "source": [
    "## `pipe.validate()` gives you the IR\n",
    "\n",
    "Before debugging *runtime* behavior, check the graph you actually built.\n",
    "`pipe.validate()` compiles the authored graph to its Intermediate Representation:\n",
    "typed nodes, edges (with sync adapter and queue size), and the execution\n",
    "topology. This is the same IR the backends and `pipe.visualize()` consume."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2b68b092",
   "metadata": {},
   "outputs": [],
   "source": [
    "ir = pipe.validate()\n",
    "\n",
    "print(f\"IR '{ir.metadata.name}' (v{ir.version}): \"\n",
    "      f\"{len(ir.nodes)} nodes, {len(ir.edges)} edges\")\n",
    "\n",
    "print(\"nodes:\")\n",
    "for node in ir.nodes:\n",
    "    print(f\"  {node.id:8s} <- {node.type}\")\n",
    "\n",
    "print(\"edges:\")\n",
    "for edge in ir.edges:\n",
    "    src = f\"{edge.source.node}.{edge.source.port}\"\n",
    "    dst = f\"{edge.destination.node}.{edge.destination.port}\"\n",
    "    print(f\"  {src:16s} -> {dst:16s} [{edge.adapter}, qsize={edge.qsize}]\")\n",
    "\n",
    "print(\"topology:\")\n",
    "print(\"  sources:\", ir.topology.sources)\n",
    "print(\"  sinks  :\", ir.topology.sinks)\n",
    "print(\"  groups :\", ir.topology.groups)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "53e8e57d",
   "metadata": {},
   "source": [
    "If the IR is wrong — a missing edge, an unexpected adapter, the wrong execution\n",
    "order — fix the wiring first. Do not start by debugging backend scheduling."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "603ec349",
   "metadata": {},
   "source": [
    "## Record a run, then replay it\n",
    "\n",
    "`pipe.record(node, path, steps=...)` steps the graph and saves a node's output\n",
    "stream to disk. A `.pkl.gz` artifact is pure `gzip` + `pickle` — no extra\n",
    "dependencies — so it round-trips anywhere `retriever-core` runs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e272ca59",
   "metadata": {},
   "outputs": [],
   "source": [
    "import tempfile\n",
    "from pathlib import Path\n",
    "\n",
    "rec_path = Path(tempfile.gettempdir()) / \"retriever_counter_stream.pkl.gz\"\n",
    "\n",
    "pipe.reset()\n",
    "buffer = pipe.record(counter, str(rec_path), steps=4, dt=0.1)\n",
    "\n",
    "print(\"recorded \", [item.value for _ts, item in buffer], \"to\", rec_path.name)\n",
    "print(\"exists  :\", rec_path.exists())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6d9dc676",
   "metadata": {},
   "source": [
    "`pipe.replay(node, path=...)` swaps that node for an in-process replay source\n",
    "that re-emits the recorded values in order. The rest of the graph is untouched,\n",
    "so downstream Flows re-run against a fixed input — no live source needed, and\n",
    "every run is identical."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e6fe690e",
   "metadata": {},
   "outputs": [],
   "source": [
    "pipe.replay(counter, path=str(rec_path))  # Counter -> recorded stream\n",
    "\n",
    "replayed = []\n",
    "for _ in range(4):\n",
    "    result = pipe.step(dt=0.1)\n",
    "    replayed.append(result.outputs[\"Doubler\"].value)\n",
    "\n",
    "pipe.close_stepper()  # finalize flows created for stepping\n",
    "\n",
    "print(\"replayed source  :\", [item.value for _ts, item in buffer])\n",
    "print(\"doubler output   :\", replayed)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "66c982fb",
   "metadata": {},
   "source": [
    "The replayed source reproduces the recorded values exactly, and `Doubler` doubles\n",
    "each one deterministically. That is the debug loop: **inspect** the IR,\n",
    "**step** in-process with a debugger, then **record once and replay many times**\n",
    "to turn a timing-sensitive run into a stable artifact you can share."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}