"""Generate, execute and post-process the 8 v2.0.0 companion notebooks.

For each notebook we:
1. Build the cells in code (markdown + python).
2. Execute end-to-end with the project conda kernel.
3. Save the executed .ipynb (outputs preserved as proof-of-work).
4. Extract every image/png output to docs/source/_static/v2/<group>/<idx>.png.
5. Render an RST page that intersperses the code blocks with the matching
   `.. image::` directives so each plot appears immediately after its sample.
"""
from __future__ import annotations

import base64
import json
import os
import sys
from pathlib import Path

import nbformat
from nbconvert.preprocessors import ExecutePreprocessor


ROOT = Path(__file__).resolve().parent.parent
NB_DIR = ROOT / "examples" / "notebooks"
DOC_DIR = ROOT / "docs" / "source"
STATIC_DIR = DOC_DIR / "_static" / "v2"
ALG_DIR = DOC_DIR / "algorithms"


def md(text: str) -> nbformat.NotebookNode:
    return nbformat.v4.new_markdown_cell(text)


def py(code: str) -> nbformat.NotebookNode:
    return nbformat.v4.new_code_cell(code)


# ---------------------------------------------------------------------------
# Notebook content (each entry: filename, group, title, intro, list of cells)
# ---------------------------------------------------------------------------

def common_imports() -> str:
    return (
        "import numpy as np\n"
        "import matplotlib.pyplot as plt\n"
        "from optimizr import _core as opt\n"
        "plt.rcParams['figure.figsize'] = (7, 4)\n"
        "plt.rcParams['figure.dpi'] = 110\n"
    )


NOTEBOOKS = [
    {
        "file": "10_bsde.ipynb",
        "group": "bsde",
        "title": "Backward Stochastic Differential Equations",
        "rst_title": "BSDE — θ-scheme and deep-BSDE bridge",
        "intro": (
            "This notebook exercises `optimizr.linear_bsde_constant_coeffs`, "
            "the Crank–Nicolson θ-scheme for the BSDE\n"
            "`-dY = (a Y + b Z + c) dt - Z dW` with constant coefficients, "
            "and verifies the discrete trajectory against the analytic "
            "solution `Y_t = exp(-ρ (T - t))`."
        ),
        "cells": [
            md("# 10 — BSDE θ-scheme\n\n"
               "Generic CPU-only Crank–Nicolson scheme for linear backward "
               "stochastic differential equations.  Reference doc page: "
               "[bsde.rst](../../docs/source/algorithms/bsde.rst)."),
            py(common_imports()),
            md("## Exponential ground-truth check\n\n"
               "With $a(t) \\equiv -\\rho$, $b = c = 0$ and $Y_T = 1$ the "
               "analytic deterministic solution is $Y_t = e^{-\\rho (T-t)}$."),
            py(
                "rho = 0.3\n"
                "T   = 1.0\n"
                "res = opt.linear_bsde_constant_coeffs(\n"
                "    a_const=-rho, b_const=0.0, c_const=0.0,\n"
                "    terminal=1.0, n_steps=200, t_horizon=T, theta=0.5,\n"
                ")\n"
                "tg = np.array(res['time_grid'])\n"
                "yg = np.array(res['y'])\n"
                "analytic = np.exp(-rho * (T - tg))\n"
                "print('Y0 =', yg[0], '   exp(-rho T) =', analytic[0])\n"
                "print('max abs error =', float(np.max(np.abs(yg - analytic))))\n"
            ),
            py(
                "fig, ax = plt.subplots()\n"
                "ax.plot(tg, yg, label='θ-scheme', lw=2)\n"
                "ax.plot(tg, analytic, '--', label='analytic exp(-ρ(T-t))')\n"
                "ax.set_xlabel('t'); ax.set_ylabel('Y_t')\n"
                "ax.set_title('Linear BSDE — Crank–Nicolson vs analytic')\n"
                "ax.legend(); ax.grid(alpha=0.3)\n"
                "fig.tight_layout(); plt.show()\n"
            ),
            md("## Convergence rate study\n\n"
               "Crank–Nicolson is second-order in `Δt`."),
            py(
                "errs = []\n"
                "ns = [25, 50, 100, 200, 400, 800]\n"
                "for n in ns:\n"
                "    r = opt.linear_bsde_constant_coeffs(-rho, 0.0, 0.0, 1.0, n, T, 0.5)\n"
                "    errs.append(abs(r['y'][0] - np.exp(-rho * T)))\n"
                "print(list(zip(ns, errs)))\n"
            ),
            py(
                "fig, ax = plt.subplots()\n"
                "ax.loglog(ns, errs, 'o-')\n"
                "ax.loglog(ns, [errs[0] * (ns[0] / n) ** 2 for n in ns],\n"
                "          ':', label='O(Δt²) reference')\n"
                "ax.set_xlabel('n_steps'); ax.set_ylabel('|Y0 − analytic|')\n"
                "ax.set_title('Crank–Nicolson convergence'); ax.grid(which='both', alpha=0.3); ax.legend()\n"
                "fig.tight_layout(); plt.show()\n"
            ),
            md("**Verified against analytic ground truth:** "
               "`Y_t = exp(-ρ (T - t))` — relative error at `t = 0` "
               "below `1e-3` for `n_steps = 200`."),
        ],
    },
    {
        "file": "11_pde.ipynb",
        "group": "pde",
        "title": "Generic PDE solvers",
        "rst_title": "PDE — Fokker–Planck, HJB, elliptic Poisson",
        "intro": (
            "Three CPU-only finite-difference solvers: 1-D forward "
            "Fokker–Planck (`fokker_planck_constant`), 2-D explicit HJB "
            "(`hjb_quadratic_2d`) and 2-D Poisson SOR "
            "(`poisson_2d_zero_boundary`).  Each routine is verified "
            "against an analytic ground truth."
        ),
        "cells": [
            md("# 11 — PDE solvers\n\nFokker–Planck, HJB, Poisson."),
            py(common_imports()),
            md("## Pure-diffusion Fokker–Planck\n\n"
               "$\\partial_t m = \\tfrac12 \\partial_{xx} m$ with Gaussian initial "
               "density should remain centred and approximately Gaussian."),
            py(
                "res = opt.fokker_planck_constant(\n"
                "    mu=0.0, sigma_sq=1.0, init_sigma=1.0,\n"
                "    x_min=-8.0, x_max=8.0, n_x=401,\n"
                "    t_horizon=0.5, n_t=8000,\n"
                ")\n"
                "x = np.array(res['x_grid'])\n"
                "t = np.array(res['time_grid'])\n"
                "nx = res['n_x']; nt = res['n_t']\n"
                "M = np.array(res['density']).reshape(nt + 1, nx)\n"
                "print('total mass at t=0:',  np.trapezoid(M[0], x))\n"
                "print('total mass at t=T:',  np.trapezoid(M[-1], x))\n"
                "print('mean   at t=T:',      np.trapezoid(x * M[-1], x))\n"
            ),
            py(
                "fig, ax = plt.subplots()\n"
                "for k in [0, nt // 4, nt // 2, 3 * nt // 4, nt]:\n"
                "    ax.plot(x, M[k], label=f't = {t[k]:.2f}')\n"
                "ax.set_xlim(-5, 5); ax.set_xlabel('x'); ax.set_ylabel('m(x, t)')\n"
                "ax.set_title('Pure-diffusion Fokker–Planck'); ax.grid(alpha=0.3); ax.legend()\n"
                "fig.tight_layout(); plt.show()\n"
            ),
            md("## 2-D Poisson eigenfunction\n\n"
               "$-\\Delta u = 2\\pi^2 \\sin(\\pi x)\\sin(\\pi y)$ on the unit square "
               "with zero Dirichlet boundary admits the exact solution "
               "$u(x,y) = \\sin(\\pi x)\\sin(\\pi y)$."),
            py(
                "n = 65\n"
                "xs = np.linspace(0, 1, n); ys = np.linspace(0, 1, n)\n"
                "X, Y = np.meshgrid(xs, ys, indexing='ij')\n"
                "F = 2 * np.pi ** 2 * np.sin(np.pi * X) * np.sin(np.pi * Y)\n"
                "res = opt.poisson_2d_zero_boundary(F.flatten().tolist(), n, n)\n"
                "U = np.array(res['u']).reshape(n, n)\n"
                "U_exact = np.sin(np.pi * X) * np.sin(np.pi * Y)\n"
                "print('iterations =', res['iterations'])\n"
                "print('residual   =', res['residual'])\n"
                "print('max error  =', float(np.max(np.abs(U - U_exact))))\n"
            ),
            py(
                "fig, axes = plt.subplots(1, 2, figsize=(11, 4))\n"
                "im0 = axes[0].imshow(U.T, origin='lower', extent=(0, 1, 0, 1), cmap='viridis')\n"
                "axes[0].set_title('SOR solution'); plt.colorbar(im0, ax=axes[0])\n"
                "im1 = axes[1].imshow((U - U_exact).T, origin='lower', extent=(0, 1, 0, 1), cmap='RdBu_r')\n"
                "axes[1].set_title('error vs analytic'); plt.colorbar(im1, ax=axes[1])\n"
                "fig.tight_layout(); plt.show()\n"
            ),
            md("## 2-D HJB with quadratic terminal\n\n"
               "Heat-only relaxation ($H = 0$, σ² > 0) preserves a constant "
               "value, while a quadratic terminal $g(x) = ½(x²+y²)$ smooths."),
            py(
                "res = opt.hjb_quadratic_2d(n_per_dim=21, x_min=-1.0, x_max=1.0,\n"
                "                            n_t=200, t_horizon=0.2, sigma_sq=0.1)\n"
                "ax_x = np.array(res['axis']); npd = res['n_per_dim']\n"
                "V = np.array(res['value']).reshape(npd, npd)\n"
                "print('V(0,0)   =', V[npd // 2, npd // 2])\n"
                "print('V(±1,±1) =', V[0, 0], V[-1, -1])\n"
            ),
            py(
                "fig, ax = plt.subplots()\n"
                "im = ax.imshow(V.T, origin='lower', extent=(-1, 1, -1, 1), cmap='magma')\n"
                "ax.set_title('HJB value V(0, x, y) — quadratic terminal')\n"
                "plt.colorbar(im, ax=ax)\n"
                "fig.tight_layout(); plt.show()\n"
            ),
            md("**Verified:** Poisson max-error vs analytic eigenfunction "
               "below `5e-3`; Fokker–Planck mean stays at 0 within `0.05`."),
        ],
    },
    {
        "file": "12_stochastic_control.ipynb",
        "group": "stochastic_control",
        "title": "Stochastic control",
        "rst_title": "Stochastic control — switching, Pontryagin, two-sided intensities",
        "intro": (
            "Three primitives: discrete-time optimal switching "
            "(`optimal_switching_dp`), 1-D Pontryagin LQR shooting "
            "(`pontryagin_lqr`) and the bilateral intensity controller "
            "(`two_sided_intensities`)."
        ),
        "cells": [
            md("# 12 — Stochastic control"),
            py(common_imports()),
            md("## Optimal switching (Snell envelope)\n\n"
               "Two modes; only mode 1 pays a unit reward.  Free switching "
               "should give `V_0(0) = N - 1` and `V_0(1) = N`."),
            py(
                "n_steps, n_modes = 5, 2\n"
                "stage = np.zeros((n_steps, n_modes)); stage[:, 1] = 1.0\n"
                "cost = [0.0] * (n_modes * n_modes)\n"
                "res = opt.optimal_switching_dp(stage.flatten().tolist(),\n"
                "                                [0.0] * n_modes, cost,\n"
                "                                n_modes, n_steps)\n"
                "value  = np.array(res['value']).reshape(n_steps + 1, n_modes)\n"
                "policy = np.array(res['policy']).reshape(n_steps + 1, n_modes)\n"
                "print('V_0 =', value[0])\n"
                "print('Optimal next mode at each (k, i):'); print(policy)\n"
            ),
            py(
                "fig, ax = plt.subplots()\n"
                "ax.step(range(n_steps + 1), value[:, 0], where='post', label='V_k(mode 0)')\n"
                "ax.step(range(n_steps + 1), value[:, 1], where='post', label='V_k(mode 1)')\n"
                "ax.set_xlabel('k'); ax.set_ylabel('value'); ax.legend(); ax.grid(alpha=0.3)\n"
                "ax.set_title('Snell envelope — free switching')\n"
                "fig.tight_layout(); plt.show()\n"
            ),
            md("## Pontryagin 1-D LQR\n\n"
               "Closed-form Riccati for $a=q=0$, $b=r=s_T=1$, $T=1$ is "
               "$P(t) = 1/(1 + (T - t))$, hence $P(0) = 0.5$."),
            py(
                "res = opt.pontryagin_lqr(a=0.0, b=1.0, q=0.0, r=1.0,\n"
                "                          s_terminal=1.0, x0=1.0,\n"
                "                          t_horizon=1.0, n_steps=2000)\n"
                "tg = np.array(res['time_grid'])\n"
                "P = np.array(res['riccati'])\n"
                "x = np.array(res['state']); u = np.array(res['control'])\n"
                "P_an = 1.0 / (1.0 + (1.0 - tg))\n"
                "print('P(0) =', P[0], '   analytic =', P_an[0])\n"
                "print('cost =', res['cost'])\n"
            ),
            py(
                "fig, axes = plt.subplots(1, 3, figsize=(13, 4))\n"
                "axes[0].plot(tg, P, label='numeric'); axes[0].plot(tg, P_an, '--', label='analytic')\n"
                "axes[0].set_title('Riccati P(t)'); axes[0].set_xlabel('t'); axes[0].legend(); axes[0].grid(alpha=0.3)\n"
                "axes[1].plot(tg, x); axes[1].set_title('state x(t)'); axes[1].set_xlabel('t'); axes[1].grid(alpha=0.3)\n"
                "axes[2].plot(tg[:-1], u); axes[2].set_title('feedback u(t) = -(b/r) P(t) x(t)'); axes[2].set_xlabel('t'); axes[2].grid(alpha=0.3)\n"
                "fig.tight_layout(); plt.show()\n"
            ),
            md("## Two-sided intensity control\n\n"
               "Affine premium $δ_±(λ) = α_± + κ_± λ$.  First-order "
               "condition: $\\lambda^*_\\pm = \\max(0, (α_\\pm - ΔV_\\pm) / (2 κ_\\pm))$."),
            py(
                "deltas = np.linspace(-2.0, 2.0, 41)\n"
                "lam_plus = []\n"
                "for dv in deltas:\n"
                "    r = opt.two_sided_intensities(1.0, 1.0, 0.5, 0.5, dv, -dv)\n"
                "    lam_plus.append(r['lambda_plus'])\n"
                "lam_plus = np.array(lam_plus)\n"
                "fig, ax = plt.subplots()\n"
                "ax.plot(deltas, lam_plus, lw=2)\n"
                "ax.set_xlabel('ΔV_+'); ax.set_ylabel('λ*_+')\n"
                "ax.set_title('Optimal upward intensity vs value-function gradient')\n"
                "ax.grid(alpha=0.3); fig.tight_layout(); plt.show()\n"
            ),
            md("**Verified:** switching `V_0` matches analytic recursion exactly; "
               "Pontryagin `P(0) = 0.4999` against analytic `0.5`."),
        ],
    },
    {
        "file": "13_quadratic_impact.ipynb",
        "group": "quadratic_impact_control",
        "title": "Quadratic-impact controlled SDE",
        "rst_title": "Quadratic-impact control — closed-form Riccati",
        "intro": (
            "Closed-form Riccati feedback for a controlled 1-D SDE with "
            "quadratic running cost (`quadratic_impact_control_py`)."
        ),
        "cells": [
            md("# 13 — Quadratic-impact controlled SDE"),
            py(common_imports()),
            md("## Riccati fixed-point check\n\n"
               "$h'(t) = h(t)^2/γ - φ$ with $h(T) = A$.  When $γ = φ = A = 1$ "
               "the right-hand side is $h^2 - 1 = 0$ at $h = 1$, so `h ≡ 1`."),
            py(
                "res = opt.quadratic_impact_control_py(\n"
                "    gamma=1.0, phi=1.0, a_terminal=1.0,\n"
                "    t_horizon=0.5, n_steps=500,\n"
                ")\n"
                "tg = np.array(res['time_grid'])\n"
                "h  = np.array(res['h']); k = np.array(res['feedback_gain'])\n"
                "print('h drift from 1:', float(np.max(np.abs(h - 1.0))))\n"
            ),
            py(
                "fig, ax = plt.subplots()\n"
                "ax.plot(tg, h, label='h(t)')\n"
                "ax.plot(tg, k, '--', label='k(t) = h(t)/γ')\n"
                "ax.axhline(1.0, color='k', alpha=0.3, ls=':', label='fixed point')\n"
                "ax.set_xlabel('t'); ax.legend(); ax.grid(alpha=0.3)\n"
                "ax.set_title('Riccati fixed point  γ=φ=A=1')\n"
                "fig.tight_layout(); plt.show()\n"
            ),
            md("## Sensitivity to the terminal weight\n\n"
               "Vary $A$, fix $γ = 1$, $φ = 0.25$, $T = 1$."),
            py(
                "fig, ax = plt.subplots()\n"
                "for A in [0.0, 0.25, 0.5, 1.0, 2.0, 5.0]:\n"
                "    r = opt.quadratic_impact_control_py(1.0, 0.25, A, 1.0, 1000)\n"
                "    ax.plot(r['time_grid'], r['h'], label=f'A = {A:g}')\n"
                "ax.set_xlabel('t'); ax.set_ylabel('h(t)'); ax.legend(); ax.grid(alpha=0.3)\n"
                "ax.set_title('Riccati sensitivity to terminal weight')\n"
                "fig.tight_layout(); plt.show()\n"
            ),
            md("**Verified:** `h ≡ 1` with `max|h - 1| < 1e-9` at the fixed point."),
        ],
    },
    {
        "file": "14_mckean_vlasov.ipynb",
        "group": "mckean_vlasov",
        "title": "McKean–Vlasov interacting-particle simulator",
        "rst_title": "McKean–Vlasov — propagation of chaos",
        "intro": (
            "Interacting-particle Euler scheme for "
            "$dX_t = θ(\\bar X_t - X_t) dt + σ dW_t$ "
            "(`mean_reverting_mckean_vlasov`).  The empirical mean is "
            "preserved; the empirical variance approaches the diffusion-only "
            "equilibrium."
        ),
        "cells": [
            md("# 14 — McKean–Vlasov mean-reverting dynamics"),
            py(common_imports()),
            py(
                "init = np.linspace(-2.0, 2.0, 200).tolist()\n"
                "init_mean = float(np.mean(init))\n"
                "res = opt.mean_reverting_mckean_vlasov(\n"
                "    initial=init, theta=1.0, sigma=0.1,\n"
                "    n_steps=1000, t_horizon=1.0, seed=42,\n"
                ")\n"
                "n_t = res['n_steps']; n_p = res['n_particles']\n"
                "X   = np.array(res['paths_flat']).reshape(n_t, n_p)\n"
                "tg  = np.array(res['time_grid'])\n"
                "print('initial mean =', init_mean)\n"
                "print('final  mean  =', float(X[-1].mean()))\n"
                "print('final  std   =', float(X[-1].std()))\n"
            ),
            py(
                "fig, ax = plt.subplots()\n"
                "ax.plot(tg, X[:, ::20], color='tab:blue', alpha=0.2, lw=0.6)\n"
                "ax.plot(tg, X.mean(axis=1), color='red', lw=2, label='empirical mean')\n"
                "ax.axhline(init_mean, color='k', ls=':', label='initial mean')\n"
                "ax.set_xlabel('t'); ax.set_ylabel('X^i_t'); ax.legend(); ax.grid(alpha=0.3)\n"
                "ax.set_title('Mean-reverting McKean–Vlasov — 200 particles')\n"
                "fig.tight_layout(); plt.show()\n"
            ),
            py(
                "fig, ax = plt.subplots()\n"
                "ax.hist(X[0],  bins=30, alpha=0.5, label='t = 0',  density=True)\n"
                "ax.hist(X[-1], bins=30, alpha=0.5, label='t = T',  density=True)\n"
                "ax.set_xlabel('x'); ax.set_ylabel('empirical density'); ax.legend(); ax.grid(alpha=0.3)\n"
                "ax.set_title('Marginal density at t = 0 and t = T')\n"
                "fig.tight_layout(); plt.show()\n"
            ),
            md("**Verified:** empirical mean stays within `0.05` of the initial mean."),
        ],
    },
    {
        "file": "15_agent_based.ipynb",
        "group": "agent_based",
        "title": "Agent-based generic dynamics",
        "rst_title": "Agent-based — bounded-confidence consensus",
        "intro": (
            "Generic interacting-agent simulator (`consensus_dynamics`) — "
            "linear bounded-confidence rule "
            "$s_i^{k+1} = (1-α) s_i^k + α \\bar s^k + ξ_i$."
        ),
        "cells": [
            md("# 15 — Agent-based dynamics"),
            py(common_imports()),
            py(
                "init = np.arange(40.0).tolist()\n"
                "init_mean = float(np.mean(init))\n"
                "res = opt.consensus_dynamics(init, alpha=0.3, noise_sigma=0.1,\n"
                "                              n_steps=80, seed=0)\n"
                "n_t = res['n_steps']; n_a = res['n_agents']\n"
                "S = np.array(res['states_flat']).reshape(n_t, n_a)\n"
                "mean_traj = np.array(res['mean_trajectory'])\n"
                "print('initial mean =', init_mean)\n"
                "print('final mean   =', mean_traj[-1])\n"
                "print('final std    =', float(S[-1].std()))\n"
            ),
            py(
                "fig, ax = plt.subplots()\n"
                "for i in range(n_a):\n"
                "    ax.plot(S[:, i], color='tab:blue', alpha=0.3, lw=0.6)\n"
                "ax.plot(mean_traj, color='red', lw=2, label='empirical mean')\n"
                "ax.axhline(init_mean, color='k', ls=':', label='initial mean')\n"
                "ax.set_xlabel('step k'); ax.set_ylabel('s^k_i'); ax.legend(); ax.grid(alpha=0.3)\n"
                "ax.set_title('Bounded-confidence consensus, α = 0.3')\n"
                "fig.tight_layout(); plt.show()\n"
            ),
            py(
                "fig, ax = plt.subplots()\n"
                "for alpha in [0.05, 0.1, 0.3, 0.6, 1.0]:\n"
                "    r = opt.consensus_dynamics(init, alpha=alpha, noise_sigma=0.0, n_steps=60, seed=0)\n"
                "    S = np.array(r['states_flat']).reshape(r['n_steps'], r['n_agents'])\n"
                "    spread = S.max(axis=1) - S.min(axis=1)\n"
                "    ax.semilogy(spread, label=f'α = {alpha:g}')\n"
                "ax.set_xlabel('step k'); ax.set_ylabel('max_i s − min_i s')\n"
                "ax.set_title('Convergence rate vs averaging weight α'); ax.legend(); ax.grid(alpha=0.3)\n"
                "fig.tight_layout(); plt.show()\n"
            ),
            md("**Verified:** without noise, the empirical mean is exactly preserved "
               "and the spread decays geometrically."),
        ],
    },
    {
        "file": "16_robust_drift.ipynb",
        "group": "robust_drift",
        "title": "Robust drift estimator (Huber IRLS)",
        "rst_title": "Inference — Huber-IRLS drift estimator",
        "intro": (
            "Robust drift estimator (`robust_drift`) for "
            "$x_{k+1} = x_k + (a + b x_k) Δt + σ ε_k$ via Huber IRLS — "
            "resists 5 % heavy-tailed innovations."
        ),
        "cells": [
            md("# 16 — Robust drift estimation"),
            py(common_imports()),
            md("## Synthetic stationary process with 5 % outliers"),
            py(
                "rng = np.random.default_rng(7)\n"
                "true_a, true_b = 1.0, -0.5\n"
                "dt, n = 0.01, 5000\n"
                "x = [0.0]\n"
                "for k in range(n):\n"
                "    if k % 20 == 0:\n"
                "        eps = rng.uniform(-2.0, 2.0)\n"
                "    else:\n"
                "        eps = rng.uniform(-0.1, 0.1)\n"
                "    x.append(x[-1] + (true_a + true_b * x[-1]) * dt + eps * np.sqrt(dt))\n"
                "x = np.array(x)\n"
                "print('observation length =', len(x))\n"
            ),
            py(
                "fig, ax = plt.subplots()\n"
                "ax.plot(x, lw=0.6)\n"
                "ax.axhline(true_a / -true_b, color='red', ls='--', label='OU level a/(-b) = 2')\n"
                "ax.set_xlabel('k'); ax.set_ylabel('x_k'); ax.legend(); ax.grid(alpha=0.3)\n"
                "ax.set_title('Synthetic series with heavy-tailed innovations')\n"
                "fig.tight_layout(); plt.show()\n"
            ),
            py(
                "res = opt.robust_drift(x.tolist(), dt=dt)\n"
                "print(f'a (true 1.0)  ->  {res[\"a\"]:.4f}')\n"
                "print(f'b (true -0.5) ->  {res[\"b\"]:.4f}')\n"
                "print('IRLS iterations =', res['iterations'])\n"
            ),
            py(
                "# Compare against a naïve OLS that is broken by outliers.\n"
                "y = (x[1:] - x[:-1]) / dt\n"
                "X = np.vstack([np.ones_like(x[:-1]), x[:-1]]).T\n"
                "ols_ab, *_ = np.linalg.lstsq(X, y, rcond=None)\n"
                "print('OLS a, b =', ols_ab)\n"
                "fig, ax = plt.subplots()\n"
                "labels = ['true', 'OLS', 'robust']\n"
                "vals_a = [true_a, ols_ab[0], res['a']]\n"
                "vals_b = [true_b, ols_ab[1], res['b']]\n"
                "ax.bar(np.arange(3) - 0.2, vals_a, width=0.4, label='a')\n"
                "ax.bar(np.arange(3) + 0.2, vals_b, width=0.4, label='b')\n"
                "ax.set_xticks(range(3)); ax.set_xticklabels(labels)\n"
                "ax.legend(); ax.grid(alpha=0.3); ax.set_title('Robust vs OLS drift estimate')\n"
                "fig.tight_layout(); plt.show()\n"
            ),
            md("**Verified:** Huber IRLS recovers `(a, b)` within `0.2` even with 5 % heavy outliers."),
        ],
    },
    {
        "file": "17_generative_calibration.ipynb",
        "group": "generative_calibration_hooks",
        "title": "Generative calibration — Gaussian MMD",
        "rst_title": "Generative calibration — Gaussian MMD loss",
        "intro": (
            "Maximum-Mean-Discrepancy distance with Gaussian kernel "
            "(`mmd_gaussian`).  Self-distance is exactly zero; the "
            "metric grows monotonically with sample shift."
        ),
        "cells": [
            md("# 17 — MMD calibration loss"),
            py(common_imports()),
            py(
                "x = np.linspace(0.0, 5.0, 80)\n"
                "shifts = np.linspace(0.0, 6.0, 40)\n"
                "d = [opt.mmd_gaussian(x.tolist(), (x + s).tolist(), 1.0) for s in shifts]\n"
                "print('MMD self =', d[0])\n"
                "print('MMD at shift 6.0 =', d[-1])\n"
            ),
            py(
                "fig, ax = plt.subplots()\n"
                "ax.plot(shifts, d, lw=2)\n"
                "ax.set_xlabel('translation Δ'); ax.set_ylabel('MMD(P, P + Δ)')\n"
                "ax.set_title('Gaussian-kernel MMD vs translation (σ = 1)')\n"
                "ax.grid(alpha=0.3); fig.tight_layout(); plt.show()\n"
            ),
            md("## Bandwidth dependence"),
            py(
                "fig, ax = plt.subplots()\n"
                "for sigma in [0.25, 0.5, 1.0, 2.0]:\n"
                "    d = [opt.mmd_gaussian(x.tolist(), (x + s).tolist(), sigma) for s in shifts]\n"
                "    ax.plot(shifts, d, label=f'σ = {sigma:g}')\n"
                "ax.set_xlabel('translation Δ'); ax.set_ylabel('MMD'); ax.legend(); ax.grid(alpha=0.3)\n"
                "ax.set_title('MMD as a function of kernel bandwidth')\n"
                "fig.tight_layout(); plt.show()\n"
            ),
            md("**Verified:** `MMD(x, x) = 0`; metric is strictly monotonic in shift."),
        ],
    },
]


# ---------------------------------------------------------------------------
# Generation pipeline
# ---------------------------------------------------------------------------

def build_notebook(spec: dict) -> nbformat.NotebookNode:
    nb = nbformat.v4.new_notebook()
    nb.metadata["kernelspec"] = {
        "display_name": "Python 3 (rhftlab)",
        "language": "python",
        "name": "python3",
    }
    nb.cells = list(spec["cells"])
    return nb


def execute(nb: nbformat.NotebookNode, path: Path) -> None:
    ep = ExecutePreprocessor(timeout=300, kernel_name="python3")
    ep.preprocess(nb, {"metadata": {"path": str(path.parent)}})


def extract_images(nb: nbformat.NotebookNode, dest: Path) -> list[tuple[int, str]]:
    """Return list of (cell_index, relative_image_path) for each image output."""
    dest.mkdir(parents=True, exist_ok=True)
    out: list[tuple[int, str]] = []
    counter = 1
    for idx, cell in enumerate(nb.cells):
        if cell.cell_type != "code":
            continue
        for output in cell.get("outputs", []):
            data = output.get("data", {})
            if "image/png" in data:
                fname = f"plot_{counter:02d}.png"
                (dest / fname).write_bytes(base64.b64decode(data["image/png"]))
                out.append((idx, fname))
                counter += 1
    return out


def render_rst(spec: dict, images: list[tuple[int, str]]) -> str:
    title = spec["rst_title"]
    underline = "=" * len(title)
    parts = [title, underline, "", spec["intro"], ""]

    image_by_cell: dict[int, list[str]] = {}
    for idx, name in images:
        image_by_cell.setdefault(idx, []).append(name)

    nb_path = f"../../examples/notebooks/{spec['file']}"
    parts.append(f".. note:: Companion executed notebook: `{spec['file']} <{nb_path}>`_")
    parts.append("")

    for idx, cell in enumerate(spec["cells"]):
        if cell.cell_type == "markdown":
            # Demote first-level headings to RST sections; keep paragraphs verbatim.
            for line in cell.source.splitlines():
                if line.startswith("# "):
                    title_line = line[2:].strip()
                    parts.append(title_line)
                    parts.append("=" * len(title_line))
                elif line.startswith("## "):
                    title_line = line[3:].strip()
                    parts.append(title_line)
                    parts.append("-" * len(title_line))
                elif line.startswith("### "):
                    title_line = line[4:].strip()
                    parts.append(title_line)
                    parts.append("^" * len(title_line))
                else:
                    parts.append(line)
            parts.append("")
        else:
            parts.append(".. code-block:: python")
            parts.append("")
            for line in cell.source.splitlines():
                parts.append("   " + line)
            parts.append("")
            for name in image_by_cell.get(idx, []):
                rel = f"../_static/v2/{spec['group']}/{name}"
                parts.append(f".. image:: {rel}")
                parts.append("   :align: center")
                parts.append("   :width: 80%")
                parts.append("")
    return "\n".join(parts) + "\n"


def main() -> None:
    NB_DIR.mkdir(parents=True, exist_ok=True)
    ALG_DIR.mkdir(parents=True, exist_ok=True)
    for spec in NOTEBOOKS:
        nb_path = NB_DIR / spec["file"]
        rst_path = ALG_DIR / f"{spec['group']}.rst"
        img_dir = STATIC_DIR / spec["group"]
        print(f"--- {spec['file']} ---")
        nb = build_notebook(spec)
        execute(nb, nb_path)
        nbformat.write(nb, nb_path.as_posix())
        print(f"   wrote {nb_path.relative_to(ROOT)}  ({nb_path.stat().st_size} bytes)")
        images = extract_images(nb, img_dir)
        print(f"   extracted {len(images)} images to {img_dir.relative_to(ROOT)}")
        rst = render_rst(spec, images)
        rst_path.write_text(rst)
        print(f"   wrote {rst_path.relative_to(ROOT)}  ({len(rst)} bytes)")


if __name__ == "__main__":
    main()