{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Operations Guide\n",
    "\n",
    "This guide provides detailed explanations and examples for each state operation available in StateTracker.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Product (Cartesian Product)\n",
    "\n",
    "The **product** operation creates a state whose states enumerate all combinations of parent state states—the Cartesian product.\n",
    "\n",
    "### Using `product()`\n",
    "\n",
    "Use the `product()` function to create a Cartesian product:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-01-17T21:21:50.061675Z",
     "iopub.status.busy": "2026-01-17T21:21:50.061486Z",
     "iopub.status.idle": "2026-01-17T21:21:50.940283Z",
     "shell.execute_reply": "2026-01-17T21:21:50.939559Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "A=0, B=0\n",
      "A=1, B=0\n",
      "A=0, B=1\n",
      "A=1, B=1\n",
      "A=0, B=2\n",
      "A=1, B=2\n"
     ]
    }
   ],
   "source": [
    "from statetracker import Manager, State, product\n",
    "\n",
    "with Manager():\n",
    "    A = State(num_values=2, name=\"A\")\n",
    "    B = State(num_values=3, name=\"B\")\n",
    "\n",
    "    C = product([A, B])  # 6 states (2 × 3)\n",
    "\n",
    "    for _ in C:\n",
    "        print(f\"A={A.value}, B={B.value}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The first state (A) varies fastest, cycling through all its states before the second state (B) advances.\n",
    "\n",
    "You can combine more than two states:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-01-17T21:21:50.962886Z",
     "iopub.status.busy": "2026-01-17T21:21:50.962588Z",
     "iopub.status.idle": "2026-01-17T21:21:50.966133Z",
     "shell.execute_reply": "2026-01-17T21:21:50.965409Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "D has 8 states\n"
     ]
    }
   ],
   "source": [
    "with Manager():\n",
    "    A = State(num_values=2, name=\"A\")\n",
    "    B = State(num_values=2, name=\"B\")\n",
    "    C = State(num_values=2, name=\"C\")\n",
    "\n",
    "    D = product([A, B, C], name=\"D\")  # 8 states\n",
    "    print(f\"D has {D.num_values} states\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Note:** `product()` does not allow duplicate states. Use `ordered_product()` if you need automatic deduplication.\n",
    "\n",
    "### Using `ordered_product()`\n",
    "\n",
    "The `ordered_product()` function automatically:\n",
    "\n",
    "- Removes duplicate states\n",
    "- Flattens nested products\n",
    "- Orders states by `iter_order` and `id`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-01-17T21:21:50.967938Z",
     "iopub.status.busy": "2026-01-17T21:21:50.967790Z",
     "iopub.status.idle": "2026-01-17T21:21:50.970778Z",
     "shell.execute_reply": "2026-01-17T21:21:50.970152Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "C has 6 states\n"
     ]
    }
   ],
   "source": [
    "from statetracker import ordered_product\n",
    "\n",
    "with Manager():\n",
    "    A = State(num_values=2, name=\"A\")\n",
    "    B = State(num_values=3, name=\"B\")\n",
    "\n",
    "    # Duplicates are automatically removed\n",
    "    C = ordered_product([A, B, A])  # Same as ordered_product([A, B])\n",
    "    print(f\"C has {C.num_values} states\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Control the ordering with `set_product_order_mode()`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-01-17T21:21:50.972793Z",
     "iopub.status.busy": "2026-01-17T21:21:50.972625Z",
     "iopub.status.idle": "2026-01-17T21:21:50.975035Z",
     "shell.execute_reply": "2026-01-17T21:21:50.974437Z"
    }
   },
   "outputs": [],
   "source": [
    "from statetracker import set_product_order_mode\n",
    "\n",
    "set_product_order_mode(\"first_state_fastest\")  # Default behavior\n",
    "# set_product_order_mode('first_state_slowest')  # Reverse ordering"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Stack (Disjoint Union)\n",
    "\n",
    "The **stack** operation creates a state that iterates through each parent state sequentially—a disjoint union. Only one parent is active at a time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-01-17T21:21:50.976526Z",
     "iopub.status.busy": "2026-01-17T21:21:50.976411Z",
     "iopub.status.idle": "2026-01-17T21:21:50.979325Z",
     "shell.execute_reply": "2026-01-17T21:21:50.978750Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "C=0, A=0, B=None\n",
      "C=1, A=1, B=None\n",
      "C=2, A=None, B=0\n",
      "C=3, A=None, B=1\n",
      "C=4, A=None, B=2\n"
     ]
    }
   ],
   "source": [
    "from statetracker import stack\n",
    "\n",
    "with Manager():\n",
    "    A = State(num_values=2, name=\"A\")\n",
    "    B = State(num_values=3, name=\"B\")\n",
    "\n",
    "    C = stack([A, B])  # 5 states (2 + 3)\n",
    "\n",
    "    for state in C:\n",
    "        print(f\"C={state}, A={A.value}, B={B.value}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice that when A is active, B is `None` (inactive), and vice versa.\n",
    "\n",
    "## Synchronize\n",
    "\n",
    "The **sync** operation creates a state that keeps multiple parent states in lockstep—they all have the same state value."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-01-17T21:21:50.980943Z",
     "iopub.status.busy": "2026-01-17T21:21:50.980805Z",
     "iopub.status.idle": "2026-01-17T21:21:50.983630Z",
     "shell.execute_reply": "2026-01-17T21:21:50.982985Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "A=0, B=0\n",
      "A=1, B=1\n",
      "A=2, B=2\n",
      "A=3, B=3\n"
     ]
    }
   ],
   "source": [
    "from statetracker import sync\n",
    "\n",
    "with Manager():\n",
    "    A = State(num_values=4, name=\"A\")\n",
    "    B = State(num_values=4, name=\"B\")\n",
    "\n",
    "    C = sync([A, B])  # 4 states\n",
    "\n",
    "    for _ in C:\n",
    "        print(f\"A={A.value}, B={B.value}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Warning:** All states passed to `sync()` must have the same `num_states`.\n",
    "\n",
    "## Slice\n",
    "\n",
    "The **slice** operation selects a subset of states from a parent state, similar to Python list slicing.\n",
    "\n",
    "### Using Index Notation\n",
    "\n",
    "The most convenient way is Python's slice syntax:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-01-17T21:21:50.985138Z",
     "iopub.status.busy": "2026-01-17T21:21:50.985010Z",
     "iopub.status.idle": "2026-01-17T21:21:50.988168Z",
     "shell.execute_reply": "2026-01-17T21:21:50.987417Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "B (A[2:5]) has 3 states\n",
      "C (A[::2]) has 5 states\n",
      "D (A[::-1]) has 10 states\n",
      "E (A[3]) has 1 states\n"
     ]
    }
   ],
   "source": [
    "with Manager():\n",
    "    A = State(num_values=10, name=\"A\")\n",
    "\n",
    "    B = A[2:5]  # States 2, 3, 4\n",
    "    C = A[::2]  # Even states: 0, 2, 4, 6, 8\n",
    "    D = A[::-1]  # Reversed: 9, 8, 7, ..., 0\n",
    "    E = A[3]  # Single state: 3\n",
    "\n",
    "    print(f\"B (A[2:5]) has {B.num_values} states\")\n",
    "    print(f\"C (A[::2]) has {C.num_values} states\")\n",
    "    print(f\"D (A[::-1]) has {D.num_values} states\")\n",
    "    print(f\"E (A[3]) has {E.num_values} states\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Using `slice()`\n",
    "\n",
    "The `slice()` function provides the same functionality:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-01-17T21:21:50.989942Z",
     "iopub.status.busy": "2026-01-17T21:21:50.989809Z",
     "iopub.status.idle": "2026-01-17T21:21:50.992484Z",
     "shell.execute_reply": "2026-01-17T21:21:50.991865Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "B has 3 states\n",
      "C has 5 states\n"
     ]
    }
   ],
   "source": [
    "from statetracker import slice\n",
    "\n",
    "with Manager():\n",
    "    A = State(num_values=10, name=\"A\")\n",
    "\n",
    "    B = slice(A, start=2, stop=5)  # States 2, 3, 4\n",
    "    C = slice(A, step=2)  # Even states\n",
    "\n",
    "    print(f\"B has {B.num_values} states\")\n",
    "    print(f\"C has {C.num_values} states\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Repeat\n",
    "\n",
    "The **repeat** operation creates a state that cycles through the parent's states multiple times."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-01-17T21:21:50.994031Z",
     "iopub.status.busy": "2026-01-17T21:21:50.993910Z",
     "iopub.status.idle": "2026-01-17T21:21:50.996710Z",
     "shell.execute_reply": "2026-01-17T21:21:50.996065Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "B=0, A=0\n",
      "B=1, A=1\n",
      "B=2, A=2\n",
      "B=3, A=0\n",
      "B=4, A=1\n",
      "B=5, A=2\n"
     ]
    }
   ],
   "source": [
    "from statetracker import repeat\n",
    "\n",
    "with Manager():\n",
    "    A = State(num_values=3, name=\"A\")\n",
    "\n",
    "    B = repeat(A, times=2)  # 6 states\n",
    "\n",
    "    for state in B:\n",
    "        print(f\"B={state}, A={A.value}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Shuffle\n",
    "\n",
    "The **shuffle** operation creates a state with randomly permuted states. The shuffle is deterministic when a seed is provided."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-01-17T21:21:50.998033Z",
     "iopub.status.busy": "2026-01-17T21:21:50.997917Z",
     "iopub.status.idle": "2026-01-17T21:21:51.000779Z",
     "shell.execute_reply": "2026-01-17T21:21:51.000127Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "B=0, A=3\n",
      "B=1, A=1\n",
      "B=2, A=2\n",
      "B=3, A=4\n",
      "B=4, A=0\n"
     ]
    }
   ],
   "source": [
    "from statetracker import shuffle\n",
    "\n",
    "with Manager():\n",
    "    A = State(num_values=5, name=\"A\")\n",
    "\n",
    "    # Random permutation with seed for reproducibility\n",
    "    B = shuffle(A, seed=42)\n",
    "\n",
    "    for state in B:\n",
    "        print(f\"B={state}, A={A.value}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can also provide an explicit permutation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-01-17T21:21:51.002341Z",
     "iopub.status.busy": "2026-01-17T21:21:51.002205Z",
     "iopub.status.idle": "2026-01-17T21:21:51.004703Z",
     "shell.execute_reply": "2026-01-17T21:21:51.004091Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "B=0, A=3\n",
      "B=1, A=2\n",
      "B=2, A=1\n",
      "B=3, A=0\n"
     ]
    }
   ],
   "source": [
    "with Manager():\n",
    "    A = State(num_values=4, name=\"A\")\n",
    "\n",
    "    # Explicit permutation: reverse order\n",
    "    B = shuffle(A, permutation=[3, 2, 1, 0])\n",
    "\n",
    "    for state in B:\n",
    "        print(f\"B={state}, A={A.value}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Sample\n",
    "\n",
    "The **sample** operation creates a state with sampled states from the parent. This is useful for creating random subsets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-01-17T21:21:51.006343Z",
     "iopub.status.busy": "2026-01-17T21:21:51.006187Z",
     "iopub.status.idle": "2026-01-17T21:21:51.009016Z",
     "shell.execute_reply": "2026-01-17T21:21:51.008558Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "B (with replacement) has 10 states\n",
      "C (without replacement) has 10 states\n"
     ]
    }
   ],
   "source": [
    "from statetracker import sample\n",
    "\n",
    "with Manager():\n",
    "    A = State(num_values=100, name=\"A\")\n",
    "\n",
    "    # Sample 10 states with replacement\n",
    "    B = sample(A, num_values=10, seed=42)\n",
    "\n",
    "    # Sample 10 states without replacement\n",
    "    C = sample(A, num_values=10, seed=42, with_replacement=False)\n",
    "\n",
    "    print(f\"B (with replacement) has {B.num_values} states\")\n",
    "    print(f\"C (without replacement) has {C.num_values} states\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can also provide explicit sampled states:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-01-17T21:21:51.010387Z",
     "iopub.status.busy": "2026-01-17T21:21:51.010261Z",
     "iopub.status.idle": "2026-01-17T21:21:51.012701Z",
     "shell.execute_reply": "2026-01-17T21:21:51.012276Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "B=0, A=0\n",
      "B=1, A=2\n",
      "B=2, A=4\n",
      "B=3, A=6\n",
      "B=4, A=8\n"
     ]
    }
   ],
   "source": [
    "with Manager():\n",
    "    A = State(num_values=10, name=\"A\")\n",
    "\n",
    "    # Explicit states to sample\n",
    "    B = sample(A, sampled_states=[0, 2, 4, 6, 8])\n",
    "\n",
    "    for state in B:\n",
    "        print(f\"B={state}, A={A.value}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Split\n",
    "\n",
    "The **split** operation divides a state into multiple sub-states.\n",
    "\n",
    "### Equal Split\n",
    "\n",
    "Split into N roughly equal parts:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-01-17T21:21:51.014433Z",
     "iopub.status.busy": "2026-01-17T21:21:51.014311Z",
     "iopub.status.idle": "2026-01-17T21:21:51.017115Z",
     "shell.execute_reply": "2026-01-17T21:21:51.016512Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "B: 4 states\n",
      "C: 3 states\n",
      "D: 3 states\n"
     ]
    }
   ],
   "source": [
    "from statetracker import split\n",
    "\n",
    "with Manager():\n",
    "    A = State(num_values=10, name=\"A\")\n",
    "\n",
    "    # Split into 3 parts: sizes 4, 3, 3\n",
    "    B, C, D = split(A, 3)\n",
    "\n",
    "    print(f\"B: {B.num_values} states\")\n",
    "    print(f\"C: {C.num_values} states\")\n",
    "    print(f\"D: {D.num_values} states\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Proportional Split\n",
    "\n",
    "Split according to proportions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-01-17T21:21:51.018758Z",
     "iopub.status.busy": "2026-01-17T21:21:51.018609Z",
     "iopub.status.idle": "2026-01-17T21:21:51.021505Z",
     "shell.execute_reply": "2026-01-17T21:21:51.020698Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train: 80 states\n",
      "test: 20 states\n"
     ]
    }
   ],
   "source": [
    "with Manager():\n",
    "    A = State(num_values=100, name=\"A\")\n",
    "\n",
    "    # Split 80/20\n",
    "    train, test = split(A, [0.8, 0.2], names=[\"train\", \"test\"])\n",
    "\n",
    "    print(f\"train: {train.num_values} states\")\n",
    "    print(f\"test: {test.num_values} states\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Interleave\n",
    "\n",
    "The **interleave** operation creates a state that alternates between parent states' states in a round-robin fashion."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-01-17T21:21:51.023188Z",
     "iopub.status.busy": "2026-01-17T21:21:51.023046Z",
     "iopub.status.idle": "2026-01-17T21:21:51.026074Z",
     "shell.execute_reply": "2026-01-17T21:21:51.025442Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "C=0, A=0, B=None\n",
      "C=1, A=None, B=0\n",
      "C=2, A=1, B=None\n",
      "C=3, A=None, B=1\n",
      "C=4, A=2, B=None\n",
      "C=5, A=None, B=2\n"
     ]
    }
   ],
   "source": [
    "from statetracker import interleave\n",
    "\n",
    "with Manager():\n",
    "    A = State(num_values=3, name=\"A\")\n",
    "    B = State(num_values=3, name=\"B\")\n",
    "\n",
    "    C = interleave([A, B])  # 6 states\n",
    "\n",
    "    for state in C:\n",
    "        print(f\"C={state}, A={A.value}, B={B.value}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The states alternate: A's state 0, B's state 0, A's state 1, B's state 1, etc.\n",
    "\n",
    "> **Warning:** All states passed to `interleave()` must have the same `num_states`.\n",
    "\n",
    "## SyncedTo\n",
    "\n",
    "The **synced_to** operation creates a state that mirrors its parent exactly. This is useful for creating an alias or checkpoint in the state DAG."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-01-17T21:21:51.027627Z",
     "iopub.status.busy": "2026-01-17T21:21:51.027528Z",
     "iopub.status.idle": "2026-01-17T21:21:51.030189Z",
     "shell.execute_reply": "2026-01-17T21:21:51.029605Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "B=0, A=0\n",
      "B=1, A=1\n",
      "B=2, A=2\n",
      "B=3, A=3\n",
      "B=4, A=4\n"
     ]
    }
   ],
   "source": [
    "from statetracker import synced_to\n",
    "\n",
    "with Manager():\n",
    "    A = State(num_values=5, name=\"A\")\n",
    "\n",
    "    B = synced_to(A, name=\"B\")  # B mirrors A exactly\n",
    "\n",
    "    for state in B:\n",
    "        print(f\"B={state}, A={A.value}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Combining Operations\n",
    "\n",
    "Operations can be freely combined to create complex iteration patterns:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-01-17T21:21:51.031482Z",
     "iopub.status.busy": "2026-01-17T21:21:51.031375Z",
     "iopub.status.idle": "2026-01-17T21:21:51.034573Z",
     "shell.execute_reply": "2026-01-17T21:21:51.034020Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train set: 16 samples\n",
      "Test set: 4 samples\n"
     ]
    }
   ],
   "source": [
    "with Manager():\n",
    "    # Create base states for different conditions\n",
    "    control = State(num_values=10, name=\"control\")\n",
    "    treatment = State(num_values=10, name=\"treatment\")\n",
    "\n",
    "    # Stack them (20 states total)\n",
    "    all_samples = stack([control, treatment])\n",
    "\n",
    "    # Shuffle for randomization\n",
    "    randomized = shuffle(all_samples, seed=42)\n",
    "\n",
    "    # Split into train/test\n",
    "    train, test = split(randomized, [0.8, 0.2])\n",
    "\n",
    "    print(f\"Train set: {train.num_values} samples\")\n",
    "    print(f\"Test set: {test.num_values} samples\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.10.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}