onf-c2/tests/test_train.py

"""Tests for onf.graph.train.loop — the demonstration-graph retriever's training curriculum.

No HDF5/suite dependency: every test builds a small, synthetic, fully in-memory graph (straight-line
toy demos through onf.graph.core.nodes.NodeTable.from_demos + onf.graph.core.edges.EdgeSet.build,
mirroring tests/test_retrieve.py's fixture style) so this runs fast and standalone in .venv-core.
coarsen=1 throughout (rather than the production default) sidesteps a real landmine unrelated to
this module: NodeTable.validate requires the table-wide node phase to reach >0.99, which needs
a raw length congruent to 1 (mod coarsen) to land a coarse group of size 1 on the final raw frame
(see tests/test_retrieve.py's N_FRAMES/COARSEN comment) — coarsen=1 makes every raw
frame its own node and removes that constraint entirely, which is all these tests need.
"""

from __future__ import annotations

from pathlib import Path

import numpy as np
import pytest
import torch

from onf.config import GraphConfig
from onf.graph.core import schema
from onf.graph.core.edges import EdgeSet
from onf.graph.net.gnn import GraphRetrieverNet
from onf.graph.core.nodes import NodeTable
from onf.graph.train.types import QuerySpec, TrainGraphSpec
from onf.graph.train.loop import save_checkpoint, train_graph
from onf.graph.core.geometry import CleanlinessScorer
from onf.graph.train.data import label_smooth_targets, make_queries, mine_negatives
from onf.graph.train.metrics import evaluate

D = 4   # small synthetic joint dim -- any works, node_encode/edges don't care


# ======================================================================================================
# synthetic-graph fixtures
# ======================================================================================================
def _make_demo(task_id: int, seed: int, n: int, dim: int = D, base: np.ndarray | None = None) -> dict:
    """A straight-line synthetic demo (mirrors tests/test_retrieve.py's _make_demo): joint
    config drifts linearly along a random unit direction from an optional per-task base offset, so
    different tasks occupy well-separated regions of joint space and different demos of one task are
    still geometrically distinguishable (different random direction)."""
    rng = np.random.RandomState(seed)
    start = (base if base is not None else np.zeros(dim)) + rng.randn(dim) * 0.01
    direction = rng.randn(dim)
    direction /= np.linalg.norm(direction)
    t = np.arange(n, dtype=np.float32)
    q = (start[None, :] + 0.05 * t[:, None] * direction[None, :]).astype(np.float32)
    qdot = np.zeros_like(q)
    qdot[:-1] = q[1:] - q[:-1]
    qdot[-1] = qdot[-2] if n > 1 else 0.0
    grip = np.zeros(n, dtype=np.uint8)
    stage = np.zeros(n, dtype=np.int16)
    return {"q": q, "qdot": qdot, "grip": grip, "stage": stage, "task_id": task_id}


def _toy_graph(n_tasks: int = 2, demos_per_task: int = 6, length: int = 40, coarsen: int = 1,
               dim: int = D, seed: int = 0) -> tuple[NodeTable, EdgeSet]:
    rng = np.random.RandomState(seed)
    demos = []
    for ti in range(n_tasks):
        base = rng.randn(dim) * 2.0
        for di in range(demos_per_task):
            demos.append(_make_demo(ti, seed=1000 * ti + di + seed, n=length, dim=dim, base=base))
    nodes = NodeTable.from_demos(demos, task_names=[f"task{i}" for i in range(n_tasks)], coarsen=coarsen)
    edges = EdgeSet.build(nodes, k_sib=4, k_align=2, device="cpu")
    return nodes, edges


def _loop_demo(t: int, task_id: int, dim: int = D, seed: int = 0) -> dict:
    """An OUT-AND-BACK demo: q rises then returns EXACTLY to its start (a triangle profile), so the
    first and last raw frame are the identical config at maximally different phase (0.0 vs ~1.0) — the
    self-intersection geometry mine_negatives's bucket (b) exists to catch."""
    rng = np.random.RandomState(seed)
    direction = rng.randn(dim)
    direction /= np.linalg.norm(direction)
    idx = np.arange(t)
    profile = np.minimum(idx, t - 1 - idx).astype(np.float32)     # 0 .. peak .. 0
    q = (0.05 * profile[:, None] * direction[None, :]).astype(np.float32)
    qdot = np.zeros_like(q)
    qdot[:-1] = q[1:] - q[:-1]
    qdot[-1] = qdot[-2] if t > 1 else 0.0
    grip = np.zeros(t, dtype=np.uint8)
    stage = np.zeros(t, dtype=np.int16)
    return {"q": q, "qdot": qdot, "grip": grip, "stage": stage, "task_id": task_id}


def _selfx_graph(t: int = 21, dim: int = D, seed: int = 0) -> tuple[NodeTable, int, int, int]:
    """demo0 = the out-and-back loop (node 0 and node t-1 are the SAME config at phase 0.0 / ~1.0 —
    bucket (b)/(c) material for positive = t-1); demo1 = a straight-line demo of the SAME length,
    offset far from the origin (guarantees a bucket-(a) same-t_idx different-strand match at
    t_idx=t-1, without ever contaminating bucket (b)/(c)). Returns (nodes, positive, expected_a,
    expected_bc).
    """
    demo0 = _loop_demo(t, task_id=0, dim=dim, seed=seed)
    demo1 = _make_demo(task_id=0, seed=seed + 1, n=t, dim=dim, base=np.full(dim, 5.0))
    nodes = NodeTable.from_demos([demo0, demo1], task_names=["task0"], coarsen=1)
    positive = t - 1                 # last node of demo0: q == demo0's start, phase ~= 1.0
    expected_a = 2 * t - 1            # demo1's t_idx == t-1 node (different strand, same t_idx)
    expected_bc = 0                    # demo0's own start: q identical to positive, phase == 0.0
    return nodes, positive, expected_a, expected_bc


def _write_graph(tmp_path, nodes: NodeTable, edges: EdgeSet):
    d = tmp_path / "graph"
    d.mkdir(exist_ok=True)
    nodes.save(d)
    edges.save(d / schema.EDGES_NPZ)
    return d


class _RowStubField:
    """Deterministic per-row f(q), keyed on call order -- lets a test dictate exactly which window
    rows are "dirty" without a real ONFField (mirrors tests/test_retrieve.py's _StubField)."""

    def __init__(self, f_by_row):
        self._f = list(f_by_row)
        self._i = 0

    def f_value(self, q):
        v = self._f[self._i % len(self._f)]
        self._i += 1
        return v


class _NormStubField:
    """A field-shaped stub whose f_value is a pure function of q -- reusable indefinitely
    (unlike _RowStubField's call-order list), so it survives the many CleanlinessScorer.weights calls a
    full make_queries/evaluate pass makes."""

    def f_value(self, q):
        return float(np.linalg.norm(np.asarray(q, dtype=np.float64)))


# ======================================================================================================
# make_queries
# ======================================================================================================
def test_make_queries_shapes_dtypes():
    nodes, _edges = _toy_graph(n_tasks=2, demos_per_task=4, length=40)
    rng = np.random.RandomState(0)
    owners = np.arange(nodes.n_demos)
    out = make_queries(nodes, rng, n=25, owners=owners, spec=QuerySpec(window=8, advance=4))

    assert out["q_hist"].shape == (25, 8, nodes.dim)
    assert out["qdot_hist"].shape == (25, 8, nodes.dim)
    assert out["grip_hist"].shape == (25, 8)
    assert out["w_hist"].shape == (25, 8)
    assert out["q_hist"].dtype == np.float32
    assert out["qdot_hist"].dtype == np.float32
    assert out["w_hist"].dtype == np.float32
    assert out["tgt_node"].dtype == np.int64
    assert out["src_node"].dtype == np.int64
    assert out["src_owner"].dtype == np.int32
    assert not np.any(np.isnan(out["q_hist"]))


def test_make_queries_targets_ahead_same_demo():
    nodes, _edges = _toy_graph(n_tasks=2, demos_per_task=4, length=40)
    rng = np.random.RandomState(1)
    owners = np.arange(nodes.n_demos)
    out = make_queries(nodes, rng, n=80, owners=owners, spec=QuerySpec(window=8, advance=4))

    for i in range(80):
        src_v, tgt_v = int(out["src_node"][i]), int(out["tgt_node"][i])
        assert nodes.owner[src_v] == nodes.owner[tgt_v] == out["src_owner"][i]
        assert nodes.t_idx[tgt_v] >= nodes.t_idx[src_v]
        assert nodes.t_raw[tgt_v] >= nodes.t_raw[src_v]


def test_make_queries_restricted_to_requested_owners():
    nodes, _edges = _toy_graph(n_tasks=2, demos_per_task=5, length=40)
    rng = np.random.RandomState(2)
    owners = np.array([0, 2, 4])
    out = make_queries(nodes, rng, n=60, owners=owners, spec=QuerySpec(window=8, advance=4))

    assert set(out["src_owner"].tolist()) <= set(owners.tolist())
    assert set(nodes.owner[out["src_node"]].tolist()) <= set(owners.tolist())
    assert set(nodes.owner[out["tgt_node"]].tolist()) <= set(owners.tolist())


# ======================================================================================================
# make_queries: variable-length history + episode-start (entry) query class
# ======================================================================================================
def test_make_queries_variable_window_length_matches_edge_pad():
    """The true history length W varies 1..window per TRAVERSAL query (entry_frac=0 isolates it), and
    the padding is byte-identical to onf.graph.core.geometry.edge_pad_hist -- the single source of truth
    shared with QueryPreprocessor.build."""
    from onf.graph.core.geometry import edge_pad_hist

    nodes, _edges = _toy_graph(n_tasks=1, demos_per_task=3, length=60)
    rng = np.random.RandomState(0)
    owners = np.arange(nodes.n_demos)
    window = 8
    out = make_queries(nodes, rng, n=300, owners=owners, spec=QuerySpec(window=window, advance=4, entry_frac=0.0))

    assert np.all(out["is_clean"])
    eff_len = np.array([len(np.unique(out["q_hist"][i], axis=0)) for i in range(300)])
    assert eff_len.min() == 1, "no query drew the shortest (W=1) history"
    assert eff_len.max() == window, "no query drew the full (W=window) history"

    # padding parity: for a query whose effective length is w, edge_pad_hist(the last w raw rows, window)
    # must reproduce q_hist[i] exactly.
    for i in range(300):
        w = int(eff_len[i])
        raw_tail = out["q_hist"][i, -w:]
        padded, _ = edge_pad_hist(raw_tail, window)
        assert np.array_equal(padded, out["q_hist"][i])


def test_make_queries_entry_class_targets_same_demo_opening():
    """Entry (episode-start) queries: W ~ Uniform(1, window) REAL consecutive frames from the demo's
    own start (a constant offset added to all of them), target = advance raw frames past the window's
    OWN (unperturbed) end, on the SAME demo -- is_clean=False (provenance)."""
    nodes, _edges = _toy_graph(n_tasks=2, demos_per_task=5, length=40)
    rng = np.random.RandomState(3)
    owners = np.arange(nodes.n_demos)
    window = 8
    out = make_queries(nodes, rng, n=200, owners=owners, spec=QuerySpec(window=window, advance=4, entry_frac=1.0))

    assert not np.any(out["is_clean"])
    # window length is ALSO variable for entry queries now (no longer pinned at W=1) -- some queries
    # must draw the shortest and some the longest history, exactly like TRAVERSAL.
    eff_len = np.array([len(np.unique(out["q_hist"][i], axis=0)) for i in range(200)])
    assert eff_len.min() == 1, "no entry query drew the shortest (W=1) history"
    assert eff_len.max() == window, "no entry query drew the full (W=window) history"

    for i in range(200):
        owner = int(out["src_owner"][i])
        src_v, tgt_v = int(out["src_node"][i]), int(out["tgt_node"][i])
        lo = int(nodes.raw_ptr[owner])
        assert nodes.owner[src_v] == nodes.owner[tgt_v] == owner
        # src sits within the demo's OWN opening window raw frames (the window is drawn from t=0)
        assert lo <= nodes.t_raw[src_v] <= lo + window - 1 + schema.COARSEN
        assert nodes.t_raw[tgt_v] >= nodes.t_raw[src_v]                 # target is ahead-or-tied


def test_make_queries_entry_and_traversal_vh_distributions_overlap():
    """THE FIX FOR THE STAGE-1 SHORTCUT: before this, every entry query had vh IDENTICALLY ZERO (a
    single perturbed frame edge-padded across the whole window), so stage 1's "abstain" (traversal) vs
    "fire" (entry) classes were trivially separable on "is this window flat" alone -- nothing to do
    with competence. Adding the perturbation as a CONSTANT offset over a REAL multi-frame window (see
    module docstring, "EPISODE-START AS A TRAINED CASE") removes that shortcut: the two classes' vh
    distributions must now overlap.
    """
    nodes, _edges = _toy_graph(n_tasks=2, demos_per_task=8, length=60)
    owners = np.arange(nodes.n_demos)

    rng_e = np.random.RandomState(11)
    entry = make_queries(nodes, rng_e, n=300, owners=owners, spec=QuerySpec(window=8, advance=4, entry_frac=1.0))
    rng_t = np.random.RandomState(12)
    trav = make_queries(nodes, rng_t, n=300, owners=owners, spec=QuerySpec(window=8, advance=4, entry_frac=0.0))

    entry_norms = np.linalg.norm(entry["qdot_hist"], axis=-1).mean(axis=-1)   # [n] mean ||vh|| per query
    trav_norms = np.linalg.norm(trav["qdot_hist"], axis=-1).mean(axis=-1)

    assert entry_norms.max() > 1e-8, "entry vh collapsed back to all-zero -- the shortcut is back"

    # IQRs overlap: NOT the old regime where entry sits at a single point (0) strictly below traversal.
    e_lo, e_hi = np.percentile(entry_norms, [25, 75])
    t_lo, t_hi = np.percentile(trav_norms, [25, 75])
    assert e_lo <= t_hi and t_lo <= e_hi, (
        f"entry vh IQR=({e_lo:.4f},{e_hi:.4f}) and traversal vh IQR=({t_lo:.4f},{t_hi:.4f}) do not "
        "overlap -- window flatness is still a usable class tell"
    )


def test_make_queries_entry_perturbation_magnitude_matches_documented_range():
    """The entry perturbation radius is Uniform(ENTRY_LO_MULT*p50, ENTRY_HI_MULT*p99) of the corpus's
    OWN leave-one-demo-out entry noise floor -- measured here, not asserted against a hand-picked
    constant.

    Uses a SMALL-magnitude toy graph (unlike _toy_graph's default base*2.0 scale) so every
    demo start sits well inside the real Panda joint limits _entry_perturb clips to -- with
    _toy_graph's normal scale some synthetic demo starts already sit OUTSIDE those limits (an
    artifact of "any magnitude works" toy data, unrelated to a real robot corpus), which would make
    clipping, not sampling, determine the observed radius and defeat this test's purpose.
    """
    from onf.graph.train.drift import PANDA_JOINT_HI, PANDA_JOINT_LO
    from onf.graph.train.data import ENTRY_HI_MULT, ENTRY_LO_MULT, entry_perturbation_floor

    # Panda joint 4 (index 3) has an ASYMMETRIC range that excludes 0 ([-3.07, -0.07] rad) -- centering
    # every toy demo at each joint's OWN mid-range (rather than 0) keeps this small-amplitude synthetic
    # graph safely inside every real limit _entry_perturb clips to, so the observed radius reflects the
    # sampling distribution rather than an accidental per-joint clip.
    center = ((PANDA_JOINT_LO[:D] + PANDA_JOINT_HI[:D]) / 2.0).astype(np.float32)
    rng0 = np.random.RandomState(7)
    demos = []
    for ti in range(2):
        base = rng0.randn(D) * 0.05
        for di in range(6):
            demos.append(_make_demo(ti, seed=1000 * ti + di, n=40, dim=D, base=base))
    for d in demos:                                   # shrink the within-demo drift too (default 0.05*t)
        d["q"] = d["q"] * 0.05 + center
        d["qdot"] *= 0.05
    nodes = NodeTable.from_demos(demos, task_names=["t0", "t1"], coarsen=1)

    floor = entry_perturbation_floor(nodes)
    lo, hi = ENTRY_LO_MULT * floor["p50"], ENTRY_HI_MULT * floor["p99"]
    assert 0.0 < lo < hi

    rng = np.random.RandomState(4)
    owners = np.arange(nodes.n_demos)
    window = 8
    raw_ptr = np.asarray(nodes.raw_ptr)
    q_raw = np.asarray(nodes.q_raw, dtype=np.float64)
    n = 200
    out = make_queries(nodes, rng, n=n, owners=owners, spec=QuerySpec(window=window, advance=4, entry_frac=1.0))

    # the perturbation is now a CONSTANT offset over a REAL multi-frame window (module docstring,
    # "EPISODE-START AS A TRAINED CASE"), so recovering the sampled radius means reading the window's
    # OWN earliest row (which corresponds to raw frame lo, perturbed by delta) rather than always
    # q_hist[i, -1] -- the earliest real row sits at index window - w_use, w_use recovered the
    # same way test_make_queries_variable_window_length_matches_edge_pad does (unique-row count).
    radii = np.empty(n)
    for i in range(n):
        owner = int(out["src_owner"][i])
        q0 = q_raw[raw_ptr[owner]]
        w_use = len(np.unique(out["q_hist"][i], axis=0))
        q0_pert = out["q_hist"][i, window - w_use].astype(np.float64)
        radii[i] = np.linalg.norm(q0_pert - q0)

    assert radii.min() >= lo - 1e-6
    assert radii.max() <= hi + 1e-6
    # spans a real range, not collapsed to one corner by clipping
    assert radii.max() - radii.min() > 0.1 * (hi - lo)


# ======================================================================================================
# make_queries: ADVANCE=0 and the DRIFT query class
# ======================================================================================================
class _CountingRandomState:
    """np.random.RandomState proxy that counts rand() -- the class-mix draws -- and forwards the rest."""

    def __init__(self, seed: int):
        self._rng = np.random.RandomState(seed)
        self.n_rand = 0

    def rand(self, *args, **kwargs):
        self.n_rand += 1
        return self._rng.rand(*args, **kwargs)

    def __getattr__(self, name):
        return getattr(self._rng, name)


def test_make_queries_drift_frac_zero_is_bit_identical_noop():
    """drift_frac=0.0 must be opt-in to the point of invisibility: identical arrays at the same seed,
    and -- the part an equality check alone would not catch -- NOT ONE rng.rand() spent on the drift
    check, so the pre-DRIFT stream is reproduced draw for draw (the same property
    tests/test_entry_static.py pins for entry_static_frac)."""
    nodes, _edges = _toy_graph(n_tasks=2, demos_per_task=4, length=40)
    owners = np.arange(nodes.n_demos)
    n = 200
    spec = dict(window=8, entry_frac=1 / 3, entry_static_frac=0.15)

    off = make_queries(nodes, _CountingRandomState(0), n=n, owners=owners,
                       spec=QuerySpec(drift_frac=0.0, **spec))
    base = make_queries(nodes, np.random.RandomState(0), n=n, owners=owners, spec=QuerySpec(**spec))
    for key in base:
        assert np.array_equal(off[key], base[key]), key
    assert not off["is_drift"].any()

    # one rand() per row for the static check, one more for every row the static check declined --
    # and nothing else.
    rng_off = _CountingRandomState(0)
    off = make_queries(nodes, rng_off, n=n, owners=owners, spec=QuerySpec(drift_frac=0.0, **spec))
    n_static = int(off["is_entry_static"].sum())
    assert rng_off.n_rand == n + (n - n_static)

    # with the class on, exactly one further rand() per row that is neither static nor entry
    rng_on = _CountingRandomState(0)
    on = make_queries(nodes, rng_on, n=n, owners=owners, spec=QuerySpec(drift_frac=0.5, **spec))
    n_static, n_drift, n_trav = (int(on[k].sum()) for k in ("is_entry_static", "is_drift", "is_clean"))
    assert rng_on.n_rand == n + (n - n_static) + (n_drift + n_trav)
    assert n_drift > 0


def test_make_queries_drift_rows_are_a_clean_prefix_then_a_ramped_tail():
    """The property no other class has: a breakpoint INSIDE the window. Rows before it are the demo's
    own raw frames bit-for-bit; rows after it deviate by an offset that RAMPS -- equal increments per
    row, reaching the full radius at T -- rather than stepping, which would put the whole radius into
    one finite-difference step. Uses demos centred in the real Panda joint limits so _entry_perturb's
    clip is not what determines the offset (same reason as
    test_make_queries_entry_perturbation_magnitude_matches_documented_range)."""
    from onf.graph.core.geometry import edge_pad_hist
    from onf.graph.train.drift import PANDA_JOINT_HI, PANDA_JOINT_LO

    center = ((PANDA_JOINT_LO[:D] + PANDA_JOINT_HI[:D]) / 2.0).astype(np.float32)
    demos = [_make_demo(0, seed=di, n=60, dim=D) for di in range(6)]
    for d in demos:
        d["q"] = d["q"] * 0.05 + center
        d["qdot"] *= 0.05
    nodes = NodeTable.from_demos(demos, task_names=["t0"], coarsen=1)
    q_raw = np.asarray(nodes.q_raw, dtype=np.float64)

    window, n = 8, 200
    out = make_queries(nodes, np.random.RandomState(5), n=n, owners=np.arange(nodes.n_demos),
                       spec=QuerySpec(window=window, entry_frac=0.0, drift_frac=1.0))
    assert np.all(out["is_drift"]) and not np.any(out["is_clean"])

    for i in range(n):
        t_end = int(nodes.t_raw[out["src_node"][i]])
        w_use = len(np.unique(out["q_hist"][i], axis=0))          # every row of a drift window differs
        clean = q_raw[t_end - w_use + 1 : t_end + 1].astype(np.float32)
        ref = edge_pad_hist(clean, window)[0].astype(np.float64)
        deviation = np.linalg.norm(out["q_hist"][i].astype(np.float64) - ref, axis=1)

        breakpoint_row = int((deviation > 0).argmax())
        assert 0 < breakpoint_row <= window - 2, (i, breakpoint_row)   # >=1 clean and >=2 drifted rows
        assert np.all(deviation[:breakpoint_row] == 0.0), i
        step = np.diff(deviation[breakpoint_row - 1 :])
        assert np.all(step > 0), i
        assert np.allclose(step, step[0], rtol=0.05), (i, step)        # a ramp, not a step


def test_make_queries_advance_zero_labels_the_windows_own_end():
    """ADVANCE=0: the head answers "where am I NOW", so tgt_node IS the window's own end node for
    every class -- the deployed transition kernel owns the lookahead."""
    nodes, _edges = _toy_graph(n_tasks=2, demos_per_task=4, length=40)
    owners = np.arange(nodes.n_demos)
    spec = dict(window=8, entry_static_frac=0.15, drift_frac=0.5)
    assert schema.ADVANCE == 0

    now = make_queries(nodes, np.random.RandomState(6), n=150, owners=owners, spec=QuerySpec(**spec))
    assert np.array_equal(now["src_node"], now["tgt_node"])
    assert now["is_drift"].any() and now["is_entry_static"].any() and now["is_clean"].any()

    ahead = make_queries(nodes, np.random.RandomState(6), n=150, owners=owners,
                          spec=QuerySpec(advance=4, **spec))
    assert not np.array_equal(ahead["src_node"], ahead["tgt_node"])


# ======================================================================================================
# the abstain CONSEQUENCE label (measure_inter_demo_spacing / abstain_is_correct) and the shared
# train/deploy velocity plumbing -- see module docstring, "THE ABSTAIN NODE"
# ======================================================================================================
def test_measure_inter_demo_spacing_positive_and_requires_multiple_demos():
    """INTER_DEMO_SPACING (module docstring, 'LABEL BY CONSEQUENCE') is a strictly positive, finite,
    corpus-DERIVED scale -- never a hand-picked constant -- and, like
    entry_perturbation_floor, needs >= 2 demos for a leave-one-demo-out floor to mean anything.
    """
    from onf.graph.train.data import measure_inter_demo_spacing

    nodes, _edges = _toy_graph(n_tasks=2, demos_per_task=6, length=30)
    spacing = measure_inter_demo_spacing(nodes)
    assert np.isfinite(spacing) and spacing > 0.0

    single_demo = _make_demo(task_id=0, seed=0, n=30, dim=D)
    single_nodes = NodeTable.from_demos([single_demo], task_names=["t"], coarsen=1)
    with pytest.raises(ValueError, match=">= 2 demos"):
        measure_inter_demo_spacing(single_nodes)


def test_make_queries_abstain_is_correct_true_for_clean_traversal_windows():
    """A genuinely CLEAN traversal window's last frame IS a raw corpus frame (distance ~0 to the
    nearest node), so its CONSEQUENCE label must read abstain-correct=True regardless of the
    (irrelevant here) generator-provenance is_clean flag -- the two happen to agree for TRAVERSAL,
    which is the point: the old is_clean label was already right for this class, just wrong for
    ENTRY/drift (see module docstring, 'LABEL BY CONSEQUENCE')."""
    nodes, _edges = _toy_graph(n_tasks=2, demos_per_task=6, length=40)
    rng = np.random.RandomState(9)
    owners = np.arange(nodes.n_demos)
    out = make_queries(nodes, rng, n=150, owners=owners, spec=QuerySpec(entry_frac=0.0))

    assert "abstain_is_correct" in out
    assert out["abstain_is_correct"].dtype == bool
    assert np.mean(out["abstain_is_correct"]) > 0.95, (
        "clean traversal windows should overwhelmingly label abstain-correct=True"
    )


def test_finite_diff_vel_shared_by_make_queries_and_retrieve(tmp_path):
    """Divergence 6 (module docstring, 'SHARED TRAIN/DEPLOY PLUMBING'): the query-construction path
    (make_queries, via onf.graph.core.geometry.finite_diff_vel) and
    onf.graph.run.retrieve.GraphRetriever.retrieve's own live path must compute the IDENTICAL
    velocity for the IDENTICAL raw window. Before the shared helper existed, train.py
    edge-extrapolated row 0 while retrieve() zero-padded it, so the SAME window produced two
    DIFFERENT query embeddings (and abstain decisions) depending on whether it was seen at train time
    or live.
    """
    from onf.graph.run.retrieve import GraphRetriever
    from onf.graph.core.geometry import finite_diff_vel

    nodes, edges = _toy_graph(n_tasks=1, demos_per_task=2, length=30)
    net = _small_net(nodes)
    d = _write_graph(tmp_path, nodes, edges)
    save_checkpoint(net, nodes, edges, d / schema.HEAD_NPZ)
    cfg = GraphConfig(hidden=16, layers=2, seed_topk=32, device="cpu")
    retriever = GraphRetriever.load(d, cfg=cfg)

    window = cfg.hist
    q_window = nodes.q_raw[5 : 5 + window].astype(np.float64)      # a real, non-padded raw window
    assert q_window.shape[0] == window

    v_shared = finite_diff_vel(q_window.astype(np.float32))                        # (1) train's path
    _w = retriever.preprocessor.build(q_window, None)                              # (2) deploy's path
    q_pad, g_pad, padded = _w.q, _w.grip, _w.padded
    assert not padded
    v_deploy = finite_diff_vel(q_pad)

    assert np.array_equal(v_shared.astype(np.float64), v_deploy.astype(np.float64))

    qh = torch.as_tensor(q_pad, dtype=torch.float32)
    gh = torch.zeros(window)
    w = torch.ones(window)
    e_q_train = retriever.net.enc(qh, torch.as_tensor(v_shared, dtype=torch.float32), gh, w)
    e_q_deploy = retriever.net.enc(qh, torch.as_tensor(v_deploy, dtype=torch.float32), gh, w)
    assert torch.equal(e_q_train, e_q_deploy)

    # regression guard: the RETIRED zero-padded convention must actually DISAGREE with the shared one,
    # proving this is a real parity check and not a tautology that would pass no matter what.
    v_old_zero_pad = torch.zeros_like(qh)
    v_old_zero_pad[1:] = qh[1:] - qh[:-1]
    e_q_old = retriever.net.enc(qh, v_old_zero_pad, gh, w)
    assert not torch.allclose(e_q_train, e_q_old)


# ======================================================================================================
# cleanliness weighting: the property the whole query-pooling design rests on
# ======================================================================================================
def test_cleanliness_batch_dirtier_frames_get_lower_weight():
    window = 4
    q_hist = np.zeros((1, window, D), dtype=np.float32)
    stub = _RowStubField([0.0, 1.0, 2.0, 3.0])   # monotonically dirtier across the window
    w = CleanlinessScorer(stub).weights_batch(q_hist)

    assert w.shape == (1, window)
    assert np.all(np.diff(w[0]) < 0.0)           # strictly decreasing weight as the window gets dirtier


def test_cleanliness_batch_uniform_when_no_field():
    q_hist = np.zeros((3, 5, D), dtype=np.float32)
    w = CleanlinessScorer(None).weights_batch(q_hist)
    assert np.allclose(w, 1.0)


# ======================================================================================================
# mine_negatives
# ======================================================================================================
def test_mine_negatives_never_returns_positive():
    nodes, positive, _a, _bc = _selfx_graph()
    query = nodes.q[positive]
    for trial in range(20):
        rng = np.random.RandomState(trial)
        negs = mine_negatives(nodes, query, positive, rng, n_neg=8)
        assert positive not in negs.tolist()


def test_mine_negatives_returns_each_category_when_available():
    nodes, positive, expected_a, expected_bc = _selfx_graph()
    query = nodes.q[positive]

    seen: set[int] = set()
    for trial in range(30):
        rng = np.random.RandomState(trial)
        negs = mine_negatives(nodes, query, positive, rng, n_neg=8)
        seen.update(negs.tolist())

    assert expected_a in seen, "bucket (a): same t_idx, different strand, never sampled"
    assert expected_bc in seen, "bucket (b)/(c): near-q/far-phase + own-kNN node, never sampled"


def test_mine_negatives_falls_back_when_selfx_bucket_empty():
    """A single-demo, monotonic straight-line graph has NO other strand (bucket (a) structurally
    empty) and no self-intersection (bucket (b) empty: a straight line's nearest neighbours in q are
    its temporal neighbours, whose phase is close too). mine_negatives must still return n_neg
    negatives via the uniform-random fallback, never raise, never return the positive.
    """
    demo = _make_demo(task_id=0, seed=0, n=30, dim=D)
    nodes = NodeTable.from_demos([demo], task_names=["t"], coarsen=1)
    positive = 15
    rng = np.random.RandomState(0)

    negs = mine_negatives(nodes, nodes.q[positive], positive, rng, n_neg=6)
    assert len(negs) == 6
    assert positive not in negs.tolist()
    assert len(np.unique(negs)) == 6


# ======================================================================================================
# label_smooth_targets
# ======================================================================================================
def test_label_smooth_sums_to_one_peaks_at_positive_and_decays():
    demo = _make_demo(task_id=0, seed=0, n=30, dim=D)
    nodes = NodeTable.from_demos([demo], task_names=["t"], coarsen=1)
    positive = 15   # mid-demo, so the +/- window is not clipped by the demo boundary

    cand, w = label_smooth_targets(nodes, positive)

    assert np.isclose(float(w.sum()), 1.0, atol=1e-5)
    assert int(cand[np.argmax(w)]) == positive

    dt = np.abs(nodes.t_idx[cand].astype(np.int64) - int(nodes.t_idx[positive]))
    order = np.argsort(dt)
    assert np.all(np.diff(w[order]) <= 1e-9)     # weight non-increasing as |dt| grows


def test_label_smooth_multi_positive_spans_other_strands_of_the_same_task():
    """Track GR change 2: the candidate set is NO LONGER confined to positive's own strand -- every
    node sharing its task_id within +/- WHERE_PHASE_BAND phase is a valid answer (the old same-strand-
    only invariant this test used to check is exactly what the redesign replaces; see module docstring,
    "MULTI-POSITIVE WHERE LABELS")."""
    nodes, _edges = _toy_graph(n_tasks=2, demos_per_task=3, length=20)
    positive = 5
    cand, w = label_smooth_targets(nodes, positive)

    assert np.isclose(float(w.sum()), 1.0, atol=1e-5)
    assert int(cand[np.argmax(w)]) == positive
    assert np.all(nodes.task_id[cand] == nodes.task_id[positive])
    assert np.all(np.abs(nodes.phase[cand] - nodes.phase[positive]) <= schema.WHERE_PHASE_BAND + 1e-6)
    assert len(set(nodes.owner[cand].tolist())) > 1, (
        "candidate set should span more than one strand of the same task"
    )


def test_label_smooth_never_leaks_onto_another_task():
    """The one invariant the multi-positive redesign KEEPS: task_id never leaks -- an off-task node is
    never a valid WHERE answer no matter how close its phase happens to be."""
    nodes, _edges = _toy_graph(n_tasks=2, demos_per_task=3, length=20)
    positive = 5
    cand, _w = label_smooth_targets(nodes, positive)
    assert np.all(nodes.task_id[cand] == nodes.task_id[positive])


def test_label_smooth_movement_cost_weighting_prefers_the_closer_candidate():
    """When q_now is supplied, two same-task/same-phase candidates at different distances from
    q_now must NOT get equal weight -- the closer one (lower movement cost) gets more mass, and
    dominates a same-magnitude farther peer even when both sit exactly as close to positive in
    phase."""
    nodes, _edges = _toy_graph(n_tasks=1, demos_per_task=4, length=30)
    positive = 15   # owner 0
    # a same-task/phase-band peer node on a different strand, and its own q
    cand0, _w0 = label_smooth_targets(nodes, positive)
    peers = [v for v in cand0 if v != positive]
    assert peers, "need at least one cross-strand peer for this test to mean anything"
    peer = peers[0]

    q_now_near_peer = nodes.q[peer].astype(np.float64)   # q_now coincides with the peer -> peer cost 0
    cand, w = label_smooth_targets(nodes, positive, q_now_near_peer)
    w_by_node = dict(zip(cand.tolist(), w.tolist()))
    # positive still wins (true_bonus), but the peer must clearly outweigh other, farther candidates
    assert w_by_node[int(peer)] > 0.0
    other_far = [v for v in cand if v not in (positive, peer)]
    if other_far:
        assert w_by_node[int(peer)] >= max(w_by_node[v] for v in other_far)


def test_label_smooth_true_continuation_stays_highest_even_when_a_peer_is_closer():
    """The TRUE continuation must remain the single highest-weighted candidate (WHERE_TRUE_BONUS) even
    when q_now is chosen to make some OTHER same-task/phase candidate's raw movement cost lower than
    the true continuation's own -- module docstring, "nothing is lost where the strand IS
    identifiable"."""
    nodes, _edges = _toy_graph(n_tasks=1, demos_per_task=4, length=30)
    positive = 15
    cand0, _w0 = label_smooth_targets(nodes, positive)
    peers = [v for v in cand0 if v != positive]
    assert peers
    peer = peers[0]

    q_now_at_peer = nodes.q[peer].astype(np.float64)     # peer's OWN raw cost is 0 -- strictly closer
    cand, w = label_smooth_targets(nodes, positive, q_now_at_peer)
    assert int(cand[np.argmax(w)]) == positive


# ======================================================================================================
# train_graph
# ======================================================================================================
def test_gnn_metrics_flags_constant_abstain_decision_as_separation_failure():
    """MAKING COLLAPSE IMPOSSIBLE TO MISS (module docstring): a COLLAPSED decision (pred_abstain
    constant across the whole held-out batch -- exactly the original bug's symptom, checkpoint
    abstain_rate == 1.000 on both clean AND fire-should windows) must NOT read as healthy. Forcing
    the net's abstain head to a constant logit collapses BOTH abstain_false_fire_rate and
    abstain_recovery_recall to the SAME value (0 or 1), so the separation check
    false_fire < recall this module's behavioural test relies on genuinely FAILS on a collapsed
    net -- proof this harness would have caught the original bug rather than reading
    (0.000, 0.000) as a perfect score. Uses ENTRY (not the deleted drift-generator stage-2 arm) as
    the "should fire" class -- entry_frac=1.0 forces every query to be a perturbed ENTRY window.
    """
    nodes, edges = _toy_graph(n_tasks=2, demos_per_task=6, length=32)
    net = _small_net(nodes)
    with torch.no_grad():
        net.abstain[2].bias.data.fill_(1e6)     # abstain always wins -- forced constant collapse

    rng = np.random.RandomState(0)
    owners = np.arange(nodes.n_demos)
    clean = make_queries(nodes, rng, n=40, owners=owners, spec=QuerySpec(entry_frac=0.0))
    fire = make_queries(nodes, rng, n=40, owners=owners, spec=QuerySpec(entry_frac=1.0))

    cfg = GraphConfig(hidden=16, layers=2, seed_topk=32, device="cpu")
    clean_metrics = evaluate(net, nodes, edges, clean, cfg=cfg, device="cpu")
    fire_metrics = evaluate(net, nodes, edges, fire, cfg=cfg, device="cpu")

    assert clean_metrics["abstain_rate"] == 1.0
    assert fire_metrics["abstain_rate"] == 1.0
    # the collapse-detection contract this module relies on: separation must fail loudly on a
    # collapsed net, not read as "(0.000, 0.000) == perfect".
    assert not (clean_metrics["abstain_false_fire_rate"] < fire_metrics["abstain_recovery_recall"])
    # the raw rate alone already makes it unambiguous, independent of how the two class-conditional
    # rates happen to read.
    assert clean_metrics["abstain_rate"] == fire_metrics["abstain_rate"] == 1.0


def test_train_graph_survives_zero_reached_pathological_seeding(tmp_path):
    """Integration-level companion to
    test_query_loss_handles_zero_reached_via_fallback_anchor: the whole curriculum (stage 1, the
    only stage left) must run to completion, with finite losses throughout, even when EVERY query
    seeds zero nodes.
    """
    nodes, edges = _toy_graph(n_tasks=2, demos_per_task=4, length=28)
    d = _write_graph(tmp_path, nodes, edges)
    cfg = GraphConfig(hidden=12, layers=1, seed_topk=0, device="cpu")

    result = train_graph(
        d, TrainGraphSpec(stages=(1,), epochs=2, seed=0, device="cpu", save=False, cfg=cfg, n_query=40, n_eval=10),
    )
    for entry in result["history"]["stage1"]:
        if entry["epoch"] != "final":
            assert np.isfinite(entry["loss"])


def test_train_graph_records_inter_demo_spacing(tmp_path):
    """train_graph computes INTER_DEMO_SPACING ONCE per run and records it (module docstring, "LABEL BY
    CONSEQUENCE": "computed once per graph ... and recorded in the run's metrics")."""
    from onf.graph.train.data import measure_inter_demo_spacing

    nodes, edges = _toy_graph(n_tasks=2, demos_per_task=6, length=32)
    d = _write_graph(tmp_path, nodes, edges)
    cfg = GraphConfig(hidden=12, layers=1, seed_topk=32, device="cpu")

    result = train_graph(
        d, TrainGraphSpec(stages=(1,), epochs=1, seed=0, device="cpu", save=False, cfg=cfg, n_query=40, n_eval=10),
    )
    assert "inter_demo_spacing" in result
    assert np.isclose(result["inter_demo_spacing"], measure_inter_demo_spacing(nodes))


def test_train_graph_determinism_same_seed(tmp_path):
    """Same seed -> bit-identical final weights. Forced single-threaded: GraphRetrieverNet's message
    passing (a sibling module, onf.graph.net.gnn) uses index_add_ scatter-accumulation, whose summation
    ORDER (and hence exact floating-point result) is a known PyTorch non-determinism under multi-
    threaded CPU execution -- not a bug in this module, just a precondition for exact reproducibility
    that RandomState(seed)/torch.manual_seed(seed) alone do not cover.
    """
    prev_threads = torch.get_num_threads()
    torch.set_num_threads(1)
    try:
        nodes, edges = _toy_graph(n_tasks=2, demos_per_task=5, length=32)
        d = _write_graph(tmp_path, nodes, edges)
        cfg = GraphConfig(hidden=12, layers=2, seed_topk=32, device="cpu")

        spec = TrainGraphSpec(
            stages=(1,), epochs=2, seed=0, device="cpu", save=False, cfg=cfg, n_query=60, n_eval=10,
        )
        r1 = train_graph(d, spec)
        r2 = train_graph(d, spec)
    finally:
        torch.set_num_threads(prev_threads)

    sd1, sd2 = r1["net"].state_dict(), r2["net"].state_dict()
    assert set(sd1) == set(sd2)
    for k in sd1:
        assert torch.equal(sd1[k], sd2[k]), k


# ======================================================================================================
# checkpoint round-trip
# ======================================================================================================
def _small_net(nodes: NodeTable) -> GraphRetrieverNet:
    net = GraphRetrieverNet(dim=nodes.dim, hidden=16, layers=2, n_rel=schema.N_RELATIONS, agg="sum")
    mean, std = nodes.feature_stats()
    net.set_stats(mean, std)
    net.eval()
    return net


def test_checkpoint_roundtrip(tmp_path):
    """save_checkpoint writes a state dict GraphRetrieverNet.load_npz can read back
    bit-for-bit -- the graph_hash STAMP itself (and refusing a mismatched one) is
    onf.graph.run.retrieve.GraphRetriever.load's job, tested alongside that class."""
    nodes, edges = _toy_graph(n_tasks=2, demos_per_task=4, length=30)
    net = _small_net(nodes)

    path = save_checkpoint(net, nodes, edges, tmp_path / "g_head.npz")
    loaded = GraphRetrieverNet.load_npz(path)

    sd1, sd2 = net.state_dict(), loaded.state_dict()
    assert set(sd1) == set(sd2)
    for k in sd1:
        assert torch.allclose(sd1[k], sd2[k])


# ======================================================================================================
# evaluate
# ======================================================================================================
def test_evaluate_returns_documented_keys():
    nodes, edges = _toy_graph(n_tasks=2, demos_per_task=6, length=40)
    net = _small_net(nodes)

    rng = np.random.RandomState(0)
    owners = np.arange(nodes.n_demos)
    queries = make_queries(nodes, rng, n=15, owners=owners, spec=QuerySpec(field=_NormStubField()))

    cfg = GraphConfig(hidden=16, layers=2, seed_topk=32, device="cpu")
    out = evaluate(net, nodes, edges, queries, field=_NormStubField(), cfg=cfg, device="cpu")

    for key in (
        "top1", "top10", "mrr", "when_mae_steps", "when_mae_phase", "where_rmse", "own_task_frac",
        "move_dist_mean", "move_dist_median",
        "when_phase_argmax_corr", "when_phase_argmax_mae", "when_phase_expect_corr", "when_phase_expect_mae",
        "crossed_n", "crossed_top1", "crossed_top10", "crossed_mrr", "crossed_when_mae_steps",
        "crossed_where_rmse",
    ):
        assert key in out, key
    assert np.isfinite(out["move_dist_mean"])
    assert np.isfinite(out["when_phase_expect_mae"])


def test_evaluate_transfer_across_differently_sized_graph():
    """The inductive claim, made executable: one trained net runs unmodified on a graph built from a
    different suite with a different node count."""
    nodes_a, _edges_a = _toy_graph(n_tasks=2, demos_per_task=4, length=30, seed=1)
    nodes_b, edges_b = _toy_graph(n_tasks=3, demos_per_task=7, length=50, seed=2)
    assert len(nodes_a) != len(nodes_b)
    assert nodes_a.dim == nodes_b.dim

    net = _small_net(nodes_a)

    rng = np.random.RandomState(0)
    queries_b = make_queries(nodes_b, rng, n=10, owners=np.arange(nodes_b.n_demos))
    cfg = GraphConfig(hidden=16, layers=2, seed_topk=32, device="cpu")

    out = evaluate(net, nodes_b, edges_b, queries_b, cfg=cfg, device="cpu")
    assert "top1" in out and np.isfinite(out["top1"])


# ======================================================================================================
# CLASS BALANCE, PART 2 -- the abstain-loss inverse-frequency weighting (module docstring, "CLASS
# BALANCE"; PART 1, the primary mechanism, is onf.graph.train.drift's generator calibration -- see
# tests/test_drift.py -- and is out of scope for this module's own tests).
# ======================================================================================================
def test_query_loss_batch_abstain_weight_scales_gradient_linearly():
    """abstain_weight multiplies the margin term BEFORE it enters the total loss
    (loss_abstain = abstain_weight * softplus(...)), so the gradient it contributes to the abstain
    bias must scale EXACTLY linearly with abstain_weight -- unlike the softplus nonlinearity itself,
    a constant multiplier outside it passes straight through the chain rule. This is what makes
    inverse-frequency reweighting ( w_c = n / (2*n_c) ) an honest per-class GRADIENT rescaling, not
    an approximation. Single-row batch, since _query_loss_batch is the only training-step
    implementation left."""
    from onf.graph.train.data import _node_feature_matrix
    from onf.graph.train.loss import _query_loss_batch
    from onf.graph.train.types import GraphTensors, QueryBatch

    nodes, edges = _toy_graph(n_tasks=1, demos_per_task=2, length=30)
    net = _small_net(nodes)
    cfg = GraphConfig(hidden=16, layers=2, seed_topk=32, device="cpu")

    node_x_t = torch.as_tensor(_node_feature_matrix(nodes))
    et = edges.torch("cpu")
    src, dst, rel, log_idf = et["src"].long(), et["dst"].long(), et["rel"], et["log_idf"]
    phase_t = torch.as_tensor(nodes.phase.astype(np.float32))
    bin_idx_t = torch.clamp((phase_t * schema.NBINS_ALIGN).long(), max=schema.NBINS_ALIGN - 1)
    gt = GraphTensors(nodes=nodes, edges=edges, node_x_t=node_x_t, src=src, dst=dst, rel=rel,
                       log_idf=log_idf, phase_t=phase_t, bin_idx_t=bin_idx_t)

    window = cfg.hist
    q_hist = np.repeat(nodes.q_raw[10][None, :], window, axis=0).astype(np.float32)
    qh = torch.as_tensor(q_hist)[None]
    vh = torch.zeros_like(qh)
    gh, w = torch.zeros(1, window), torch.ones(1, window)
    tgt = np.array([int(nodes.demo_nodes(0)[5])])
    src_v = np.array([int(nodes.demo_nodes(0)[2])])

    def grad_at(weight: float, label: bool) -> float:
        net.zero_grad()
        rng = np.random.RandomState(0)   # same draw for mine_negatives -- isolates weight's own effect
        qb = QueryBatch(qh=qh, vh=vh, gh=gh, w=w, tgt=tgt, src_v=src_v,
                         abstain_is_correct=np.array([label]), abstain_weight=np.array([weight]),
                         is_entry_static=np.array([False]), is_drift=np.array([False]),
                         break_row=np.array([-1]), perturb_radius=np.array([0.0]))
        loss = _query_loss_batch(net, gt, qb, rng, cfg)
        loss.backward()
        return float(net.abstain[2].bias.grad.item())

    for label in (True, False):
        g1 = grad_at(1.0, label)
        g3 = grad_at(3.0, label)
        assert g3 == pytest.approx(3.0 * g1, rel=1e-4), (label, g1, g3)


# ======================================================================================================
# PHASE-PRIMARY OBJECTIVE (Track GR change 1) -- gradients must reach the message-passing TRUNK, not
# just the (auxiliary) node-ranking head, from the phase-CE/phase-expectation terms alone.
# ======================================================================================================
def test_query_loss_batch_phase_terms_alone_reach_message_passing_trunk():
    """Isolate the PRIMARY phase terms (cfg.phase_ce_w/cfg.phase_expect_w) from every other
    term (node-identity cfg.node_w=0, and the SRC/ABSTAIN auxiliaries zeroed via a custom
    onf.graph.train.types.LossSpec) and confirm gradient still reaches
    net.rel/net.node/net.upd -- the shared message-passing trunk
    onf.graph.net.gnn.GraphRetrieverNet.propagate uses -- not merely the readout head. If this
    failed, the phase objective would be training a decorative loss term that never actually shapes
    retrieval, silently defeating the whole redesign."""
    from onf.graph.train.data import _node_feature_matrix
    from onf.graph.train.loss import DEFAULT_LOSS_SPEC, _query_loss_batch
    from onf.graph.train.types import GraphTensors, QueryBatch
    import dataclasses

    zeroed_spec = dataclasses.replace(DEFAULT_LOSS_SPEC, src_aux_w=0.0, abstain_w=0.0)

    nodes, edges = _toy_graph(n_tasks=1, demos_per_task=2, length=30)
    net = _small_net(nodes)
    cfg = GraphConfig(hidden=16, layers=2, seed_topk=32, device="cpu", node_w=0.0,
                       phase_ce_w=1.0, phase_expect_w=1.0)

    node_x_t = torch.as_tensor(_node_feature_matrix(nodes))
    et = edges.torch("cpu")
    src, dst, rel, log_idf = et["src"].long(), et["dst"].long(), et["rel"], et["log_idf"]
    phase_t = torch.as_tensor(nodes.phase.astype(np.float32))
    bin_idx_t = torch.clamp((phase_t * schema.NBINS_ALIGN).long(), max=schema.NBINS_ALIGN - 1)
    gt = GraphTensors(nodes=nodes, edges=edges, node_x_t=node_x_t, src=src, dst=dst, rel=rel,
                       log_idf=log_idf, phase_t=phase_t, bin_idx_t=bin_idx_t)

    window = cfg.hist
    q_hist = np.repeat(nodes.q_raw[10][None, :], window, axis=0).astype(np.float32)
    qh = torch.as_tensor(q_hist)[None]
    vh = torch.zeros_like(qh)
    gh, w = torch.zeros(1, window), torch.ones(1, window)
    tgt = np.array([int(nodes.demo_nodes(0)[5])])
    src_v = np.array([int(nodes.demo_nodes(0)[2])])
    rng = np.random.RandomState(0)

    net.zero_grad()
    qb = QueryBatch(qh=qh, vh=vh, gh=gh, w=w, tgt=tgt, src_v=src_v,
                     abstain_is_correct=np.array([False]), abstain_weight=np.array([1.0]),
                     is_entry_static=np.array([False]), is_drift=np.array([False]),
                     break_row=np.array([-1]), perturb_radius=np.array([0.0]))
    loss = _query_loss_batch(net, gt, qb, rng, cfg, zeroed_spec)
    loss.backward()

    for mod_name in ("rel", "node", "upd"):
        mod = getattr(net, mod_name)
        grads = [p.grad for p in mod.parameters()]
        assert all(g is not None for g in grads), f"{mod_name}: no gradient reached at all"
        assert any(float(g.abs().sum()) > 0 for g in grads), f"{mod_name}: gradient is all-zero"


def test_stage_trainer_weights_abstain_loss_inversely_to_realised_class_frequency(monkeypatch):
    """StageTrainer (module docstring, "CLASS BALANCE", PART 2) must compute, from the query set's
    OWN realised abstain_is_correct frequency, w_c = n / (2 * n_c) per class and thread it into
    every _query_loss_batch call as its abstain_weight ARRAY -- verified here by monkeypatching
    _query_loss_batch with a spy that records the per-query (label, weight) pairs it was called
    with (zipping the batched abstain_is_correct/abstain_weight arrays back into per-query
    values), on a query set DELIBERATELY skewed 90/10 so the two classes' weights are clearly
    distinguishable (not both ~1.0, which a no-op bug could pass by accident). qbatch=1 forces one
    query per batched call, so this also doubles as a coverage check that every query is visited
    exactly once per epoch under the new batched loop, same as the old per-query one."""
    import onf.graph.train.loop as train_mod
    from onf.graph.train.loss import DEFAULT_LOSS_SPEC
    from onf.graph.train.types import EvalQueries, QueryData, StageSpec

    nodes, edges = _toy_graph(n_tasks=2, demos_per_task=6, length=32)
    net = _small_net(nodes)
    net.train()
    cfg = GraphConfig(hidden=16, layers=2, seed_topk=32, device="cpu", qbatch=1)
    opt = torch.optim.AdamW(net.parameters(), lr=1e-3)

    rng = np.random.RandomState(0)
    owners = np.arange(nodes.n_demos)
    n_q = 20
    queries = make_queries(nodes, rng, n=n_q, owners=owners, spec=QuerySpec(entry_frac=0.0))
    # force a KNOWN 90/10 abstain/fire split, independent of whatever make_queries happened to draw --
    # this test is about the WEIGHT FORMULA, not about make_queries' own realised balance.
    abstain_is_correct = np.zeros(n_q, dtype=bool)
    abstain_is_correct[: round(0.9 * n_q)] = True
    queries["abstain_is_correct"] = abstain_is_correct
    held_queries = make_queries(nodes, rng, n=10, owners=owners, spec=QuerySpec(entry_frac=0.0))

    n_abstain, n_fire = int(abstain_is_correct.sum()), int((~abstain_is_correct).sum())
    expected_w_abstain = n_q / (2.0 * n_abstain)
    expected_w_fire = n_q / (2.0 * n_fire)

    seen_weights: dict[bool, list[float]] = {True: [], False: []}
    real_query_loss_batch = train_mod._query_loss_batch

    def spy(*args, **kwargs):
        # StageTrainer calls _query_loss_batch(net, gt, qb, rng, cfg) positionally -- qb (a
        # QueryBatch) is the 3rd positional arg (index 2) and carries abstain_is_correct/
        # abstain_weight as its own fields; read it back out defensively (falls back to kwargs in
        # case a future refactor passes it by keyword instead).
        qb = args[2] if len(args) > 2 else kwargs["qb"]
        for label, weight in zip(np.asarray(qb.abstain_is_correct), np.asarray(qb.abstain_weight)):
            seen_weights[bool(label)].append(float(weight))
        return real_query_loss_batch(*args, **kwargs)

    monkeypatch.setattr(train_mod, "_query_loss_batch", spy)

    gt = train_mod._prepare_tensors(nodes, edges, "cpu")
    spec = StageSpec(epochs=1, rng=rng, dev="cpu", cfg=cfg, logger=None, stage_name="stageT", field=None, loss_spec=DEFAULT_LOSS_SPEC)
    evalq = EvalQueries(held=QueryData.from_dict(held_queries))
    train_mod.StageTrainer(net, opt, gt, train_mod.QuerySet.from_queries(queries, "cpu"), evalq, spec).run()

    assert len(seen_weights[True]) == n_abstain
    assert len(seen_weights[False]) == n_fire
    for wv in seen_weights[True]:
        assert wv == pytest.approx(expected_w_abstain)
    for wv in seen_weights[False]:
        assert wv == pytest.approx(expected_w_fire)
    # the point of inverse-frequency balancing: the MINORITY class (fire, 10%) gets the LARGER weight.
    assert expected_w_fire > expected_w_abstain


# ==================================================================================================
# STAGE 2 -- the chunk-blend target (onf.blend.kinematics, onf.blend.target) and the loss over it.
# ==================================================================================================
ARTIFACT_DIR = Path(__file__).resolve().parents[1] / "outputs/long/latest/artifacts"
skip_no_corpus = pytest.mark.skipif(
    not (ARTIFACT_DIR / "g_nodes.npz").exists(),
    reason=f"needs the recorded corpus at {ARTIFACT_DIR} to check FK against stored obs/ee_*",
)


def _real_nodes() -> NodeTable:
    return NodeTable.load(ARTIFACT_DIR)


@skip_no_corpus
def test_forward_kinematics_round_trips_the_recorded_end_effector_pose():
    """The FK model behind ee_now, checked against the poses LIBERO actually recorded.

    Three claims, because three different things could be wrong. (1) The absolute position model
    ee_pos = base_t + fk_pos + R @ tool_d holds, with tool_d the rigid grip-site offset -- fitted
    per TASK, because base_t is a property of the SCENE and not of the robot: it is
    [-0.660, 0, 0.912] on the kitchen scenes, [-0.510, 0, 0.420] on the living-room ones and
    [-0.750, 0, 0.912] on the study one, so one pooled fit lands 26 cm out. (2) The absolute
    orientation reproduces obs/ee_ori in ITS unnormalized, unwrapped convention, not merely up to
    that convention. (3) The claim NodeTable.ee_base rests on: world minus base is a PURE
    TRANSLATION, constant within a scene and large across scenes. If it were not, moving the
    reference segment into the base frame would rotate it, and the action deltas built from it
    would no longer be the world-frame deltas the policy's chunk is stated in.
    """
    pytest.importorskip("pytorch_kinematics")
    from onf.blend.ee_track import matrix_to_rotvec, rotvec_to_matrix
    from onf.blend.kinematics import PandaKinematics, matrix_to_ee_rotvec

    nodes = _real_nodes()
    kin = PandaKinematics()
    demo_task = np.array([int(nodes.task_id[nodes.demo_ptr[j]]) for j in range(nodes.n_demos)])
    # Several demos per task, not one: base_t and R @ tool_d are only separable across a spread of
    # wrist orientations, and a single demo does not span enough of them to condition the fit.
    demos = [int(j) for t in np.unique(demo_task) for j in np.flatnonzero(demo_task == t)[:5]]
    frames = np.concatenate(
        [np.arange(nodes.raw_ptr[j], nodes.raw_ptr[j + 1]) for j in demos]
    )
    pose = kin.flange(np.asarray(nodes.q_raw, dtype=np.float64)[frames])
    stored = np.asarray(nodes.ee_raw, dtype=np.float64)[frames]

    # (1) absolute position, one [I | R] least-squares fit per task.
    owner = np.searchsorted(nodes.raw_ptr, frames, side="right") - 1
    for t in np.unique(demo_task):
        rows = np.isin(owner, np.flatnonzero(demo_task == t))
        flange = pose.position[rows] - np.einsum("nij,j->ni", pose.rotation[rows], kin.tool_offset)
        design = np.concatenate(
            [np.tile(np.eye(3), (int(rows.sum()), 1, 1)), pose.rotation[rows]], axis=2
        ).reshape(-1, 6)
        target = (stored[rows, :3] - flange).reshape(-1)
        fit, *_ = np.linalg.lstsq(design, target, rcond=None)
        residual = np.linalg.norm((design @ fit - target).reshape(-1, 3), axis=1)
        assert fit[3:] == pytest.approx(kin.tool_offset, abs=5e-4), (t, fit)
        assert np.sqrt((residual ** 2).mean()) < 5e-4, (t, residual.max())

    # (2) absolute orientation, in obs/ee_ori's own convention.
    ori_err = np.linalg.norm(matrix_to_ee_rotvec(pose.rotation) - stored[:, 3:], axis=1)
    assert np.sqrt((ori_err ** 2).mean()) < 3e-3
    assert ori_err.max() < 3e-2

    # (3) world minus base, per scene. The position offset is a constant within a task and the
    # rotation between the two frames is the identity, so the offset drops out of every DIFFERENCE
    # the tracking chunk takes -- but only once both operands are on the same side of it.
    base = kin.pose(np.asarray(nodes.q_raw, dtype=np.float64)[frames])
    rot_err = np.linalg.norm(
        matrix_to_rotvec(
            rotvec_to_matrix(base[:, 3:]) @ np.swapaxes(rotvec_to_matrix(stored[:, 3:]), -1, -2)
        ),
        axis=1,
    )
    assert np.sqrt((rot_err ** 2).mean()) < 3e-3 and rot_err.max() < 3e-2
    offsets = []
    for t in np.unique(demo_task):
        rows = np.isin(owner, np.flatnonzero(demo_task == t))
        delta = stored[rows, :3] - base[rows, :3]
        assert delta.std(axis=0).max() < 5e-4, (t, delta.std(axis=0))
        offsets.append(delta.mean(axis=0))
    spread = np.linalg.norm(np.asarray(offsets)[:, None] - np.asarray(offsets)[None, :], axis=-1)
    assert spread.max() > 0.4, spread          # the teleport ee_base exists to remove


@skip_no_corpus
def test_a_corr_on_an_unperturbed_window_is_the_pose_derived_demo_track():
    """On a clean window the target must reduce to the demo's OWN pose-derived track.

    a_corr is EeTrack.chunk(demo segment at the label node, the window's current pose), and on a
    TRAVERSAL row that current pose is a frame the corpus recorded. So a_corr has to equal the same
    chunk built straight off that frame's own base-frame pose -- if it does not, ee_now has drifted
    off the frame the reference segment lives in and every row of the chunk is measured against the
    wrong origin.
    """
    pytest.importorskip("pytorch_kinematics")
    from onf.blend.ee_track import ACTION_LIMIT, ActionScale, EeTrack
    from onf.blend.kinematics import PandaKinematics
    from onf.blend.kinematics import PandaKinematics
    from onf.blend.target import ChunkTargetBuilder, PolicyActionCache

    nodes = _real_nodes()
    scale = ActionScale.from_json(ARTIFACT_DIR / "action_scale.json")
    builder = ChunkTargetBuilder(
        nodes=nodes, scale=scale, policy=PolicyActionCache.load(ARTIFACT_DIR).check_against(nodes),
        kinematics=PandaKinematics(), chunk_len=schema.SEG_K,
    )
    queries = make_queries(
        nodes, np.random.RandomState(0), 64, np.arange(nodes.n_demos), QuerySpec(entry_frac=0.0)
    )
    assert queries["perturb_radius"].max() == 0.0, "entry_frac=0 must give a wholly clean set"

    targets = builder.build(queries)
    track = EeTrack(scale, bound=ACTION_LIMIT)
    for i, (tgt, t_now) in enumerate(zip(queries["tgt_node"], queries["t_raw_now"])):
        expected = track.chunk(
            nodes.ee_segment(int(tgt), schema.SEG_K),
            np.asarray(nodes.ee_base, dtype=np.float64)[int(t_now)],
        )
        # The tolerance is float32's: reference and live pose now come from the SAME forward
        # kinematics, and ee_base is stored f32, so nothing else separates the two.
        assert targets.a_corr[i] == pytest.approx(expected, abs=1e-3), i

    # a_policy must be the frozen policy's chunk at the window's OWN frame, not the label node's.
    assert targets.a_policy.shape == (64, schema.SEG_K, 7)
    assert targets.a_policy[:, 0, :] == pytest.approx(
        np.load(ARTIFACT_DIR / "a_pi_raw.npz")["a_pi_raw"][queries["t_raw_now"]]
    )


def _query_batch(queries, qh, window, n_q):
    """One onf.graph.train.types.QueryBatch off a make_queries result, chunk fields left unset."""
    from onf.graph.train.types import QueryBatch

    return QueryBatch(
        qh=qh, vh=torch.zeros_like(qh), gh=torch.zeros(n_q, window), w=torch.ones(n_q, window),
        tgt=queries["tgt_node"], src_v=queries["src_node"],
        abstain_is_correct=queries["abstain_is_correct"], abstain_weight=np.ones(n_q),
        is_entry_static=queries["is_entry_static"], is_drift=queries["is_drift"],
        break_row=queries["break_row"], perturb_radius=queries["perturb_radius"],
    )


def test_chunk_loss_reaches_both_the_blend_weight_and_the_message_passing_trunk():
    """The whole point of the objective: mse(a_exec, a_corr) must train BOTH branches of the blend.

    alpha decides how much of the tracking chunk is executed and the retriever's posterior decides
    what that chunk is, so a gradient reaching only one of them would leave the other free. The
    trunk (net.rel/node/upd) is checked rather than the readout head alone for the same reason
    test_query_loss_batch_phase_terms_alone_reach_message_passing_trunk checks it.

    Also pins the torch tracking chunk against onf.blend.ee_track.EeTrack, the numpy one the deploy
    path runs: they are two statements of one construction and are allowed to differ only in dtype.
    """
    from onf.blend.alpha import AlphaNet
    from onf.blend.ee_track import ActionScale, EeTrack
    from onf.blend.target import track_chunk
    from onf.graph.train.loss import CHUNK, DEFAULT_LOSS_SPEC, ChunkLossComputer
    from onf.graph.train.data import _node_feature_matrix
    from onf.graph.train.types import GraphTensors
    import dataclasses

    nodes, edges = _toy_graph(n_tasks=1, demos_per_task=2, length=30)
    net = _small_net(nodes)
    cfg = GraphConfig(hidden=16, layers=2, seed_topk=32, device="cpu")
    seg_k, n_q = 4, 3
    scale = ActionScale(pos=0.0116, rot=0.1092, r2_pos=1.0, r2_rot=1.0)

    rng = np.random.RandomState(0)
    ee_table = rng.randn(len(nodes), seg_k, 6).astype(np.float32) * 0.1
    ee_now = rng.randn(n_q, 6).astype(np.float32) * 0.1

    # The torch and numpy tracking chunks must agree before anything is asserted about gradients,
    # both unbounded and saturated, and both on and far off the reference -- under pursuit every
    # row is a function of the rows before it, so a disagreement in one compounds through the rest.
    from onf.blend.ee_track import ACTION_LIMIT

    off_path = ee_now + np.array([0.09, -0.04, 0.02, 0.0, 0.0, 0.0], dtype=np.float32)
    for bound in (None, ACTION_LIMIT):
        for live in (ee_now, off_path):
            expected = np.stack(
                [EeTrack(scale, bound=bound).chunk(ee_table[i], live[i]) for i in range(n_q)]
            )
            got = track_chunk(torch.as_tensor(ee_table[:n_q]), torch.as_tensor(live), scale, bound)
            assert got.numpy() == pytest.approx(expected, abs=1e-5), bound
    got = track_chunk(torch.as_tensor(ee_table[:n_q]), torch.as_tensor(ee_now), scale)

    et = edges.torch("cpu")
    phase_t = torch.as_tensor(nodes.phase.astype(np.float32))
    gt = GraphTensors(
        nodes=nodes, edges=edges, node_x_t=torch.as_tensor(_node_feature_matrix(nodes)),
        src=et["src"].long(), dst=et["dst"].long(), rel=et["rel"], log_idf=et["log_idf"],
        phase_t=phase_t,
        bin_idx_t=torch.clamp((phase_t * schema.NBINS_ALIGN).long(), max=schema.NBINS_ALIGN - 1),
        ee_table=torch.as_tensor(ee_table), action_scale=scale,
    )

    window = cfg.hist
    qh = torch.as_tensor(np.repeat(nodes.q_raw[:n_q][:, None, :], window, axis=1))
    queries = make_queries(nodes, np.random.RandomState(1), n_q, np.arange(nodes.n_demos),
                            QuerySpec(entry_frac=0.0, window=window))
    qb = dataclasses.replace(
        _query_batch(queries, qh, window, n_q),
        a_policy=torch.as_tensor(rng.randn(n_q, seg_k, 7).astype(np.float32)),
        a_corr=torch.as_tensor(rng.randn(n_q, seg_k, 6).astype(np.float32)),
        ee_now=torch.as_tensor(ee_now),
    )

    alpha_net = AlphaNet(embed_dim=cfg.hidden, chunk_len=seg_k, hidden=8)
    spec = dataclasses.replace(DEFAULT_LOSS_SPEC, objective=CHUNK)
    net.zero_grad()
    alpha_net.zero_grad()
    ChunkLossComputer(net, alpha_net, cfg, spec)(gt, qb).backward()

    for name in ("rel", "node", "upd"):
        grads = [p.grad for p in getattr(net, name).parameters()]
        assert all(g is not None for g in grads), f"{name}: no gradient reached at all"
        assert any(float(g.abs().sum()) > 0 for g in grads), f"{name}: gradient is all-zero"
    # The blend weight AND the saturation amplitude, which is the row-0 magnitude bound.
    assert float(alpha_net.mlp[-1].weight.grad.abs().sum()) > 0, "no gradient reached alpha"

    # A batch missing any chunk input must say so rather than quietly training something else.
    with pytest.raises(ValueError, match="a_corr"):
        ChunkLossComputer(net, alpha_net, cfg, spec)(gt, dataclasses.replace(qb, a_corr=None))



# ======================================================================================================
# Stage 2 -- the CHUNK objective's wiring through QuerySet / StageTrainer / GraphTrainer
# ======================================================================================================
def _chunk_stage(nodes, edges, net, cfg, seg_k, n_q, n_held, seed=0):
    """One CHUNK-objective StageTrainer over a toy graph, with synthetic chunk targets.

    ChunkTargetBuilder needs a recorded corpus (ee_raw + a policy-action cache), which a toy graph
    has none of, so the targets are drawn at random here: this exercises the WIRING, and
    test_chunk_loss_reaches_both_the_blend_weight_and_the_message_passing_trunk already pins the
    numerics against onf.blend.ee_track.

    Args:
        nodes: Toy node table.
        edges: Its edge set.
        net: The retriever, already frozen by the caller.
        cfg: Graph configuration.
        seg_k: Chunk length.
        n_q: Training queries.
        n_held: Held-out queries.
        seed: Seed for the synthetic targets and the query draw.

    Returns:
        (StageTrainer, AlphaNet, optimiser).
    """
    import dataclasses

    from onf.blend.alpha import AlphaNet
    from onf.blend.ee_track import ActionScale
    from onf.blend.target import ChunkTargets
    from onf.graph.train import loop as train_mod
    from onf.graph.train.corrupt import RetrievalCorruptor
    from onf.graph.train.loss import CHUNK, DEFAULT_LOSS_SPEC
    from onf.graph.train.data import _node_feature_matrix
    from onf.graph.train.types import EvalQueries, GraphTensors, QueryData, StageSpec

    rng = np.random.RandomState(seed)
    scale = ActionScale(pos=0.0116, rot=0.1092, r2_pos=1.0, r2_rot=1.0)
    et = edges.torch("cpu")
    phase_t = torch.as_tensor(nodes.phase.astype(np.float32))
    gt = GraphTensors(
        nodes=nodes, edges=edges, node_x_t=torch.as_tensor(_node_feature_matrix(nodes)),
        src=et["src"].long(), dst=et["dst"].long(), rel=et["rel"], log_idf=et["log_idf"],
        phase_t=phase_t,
        bin_idx_t=torch.clamp((phase_t * schema.NBINS_ALIGN).long(), max=schema.NBINS_ALIGN - 1),
        ee_table=torch.as_tensor(rng.randn(len(nodes), seg_k, 6).astype(np.float32) * 0.1),
        action_scale=scale,
    )

    def make_set(n):
        queries = make_queries(nodes, rng, n, np.arange(nodes.n_demos),
                                QuerySpec(entry_frac=0.5, window=cfg.hist))
        qs = train_mod.QuerySet.from_queries(
            queries, "cpu", None, RetrievalCorruptor(nodes), rng
        )
        targets = ChunkTargets(
            a_policy=rng.randn(n, seg_k, 7).astype(np.float32) * 0.1,
            a_corr=rng.randn(n, seg_k, 6).astype(np.float32) * 0.1,
            ee_now=rng.randn(n, 6).astype(np.float32) * 0.1,
            ee_ref=rng.randn(n, seg_k, 6).astype(np.float32) * 0.1,
        )
        return queries, dataclasses.replace(qs, targets=targets)

    train_queries, train_set = make_set(n_q)
    held_queries, held_set = make_set(n_held)

    alpha_net = AlphaNet(embed_dim=cfg.hidden, chunk_len=seg_k, hidden=8)
    opt = torch.optim.AdamW(alpha_net.parameters(), lr=1e-2)
    stage = train_mod.StageTrainer(
        net=net, opt=opt, gt=gt, queries=train_set,
        evalq=EvalQueries(held=QueryData.from_dict(held_queries), held_set=held_set),
        spec=StageSpec(epochs=1, rng=rng, dev="cpu", cfg=cfg, logger=None, stage_name="stage2",
                        field=None, loss_spec=dataclasses.replace(DEFAULT_LOSS_SPEC, objective=CHUNK)),
        alpha_net=alpha_net,
    )
    return stage, alpha_net, opt


def test_chunk_stage_trains_alpha_alone_and_reports_the_gate(capsys):
    """The Stage 2 smoke: one CHUNK epoch runs, the loss is finite, ONLY alpha moves, gate prints.

    The three claims are one claim. A finite loss says the corruption, the reference segment and the
    DTW teacher all reached the batch in the shapes the loss expects; an unchanged retriever says
    the freeze held; and the stratified gate line is the number the whole stage exists to produce,
    so a run that trains without printing it cannot be judged.
    """
    from onf.graph.train.corrupt import KIND_NAMES

    nodes, edges = _toy_graph(n_tasks=2, demos_per_task=4, length=30)
    cfg = GraphConfig(hidden=16, layers=2, seed_topk=32, device="cpu", qbatch=8)
    net = _small_net(nodes)
    net.requires_grad_(False)
    net.eval()

    stage, alpha_net, _ = _chunk_stage(nodes, edges, net, cfg, seg_k=4, n_q=16, n_held=12)
    before_net = {k: v.detach().clone() for k, v in net.state_dict().items()}
    before_alpha = {k: v.detach().clone() for k, v in alpha_net.state_dict().items()}

    history = stage.run()

    assert np.isfinite(history[0]["loss"]), history[0]["loss"]
    assert np.isfinite(history[0]["chunk_mse"]), "model selection could not score the blended chunk"
    for k, v in net.state_dict().items():
        assert torch.equal(v, before_net[k]), f"the frozen retriever's {k} moved"
    assert any(
        not torch.equal(v, before_alpha[k]) for k, v in alpha_net.state_dict().items()
    ), "one epoch left the blend weight untouched"

    strata = stage.report_alpha_strata()
    assert sum(strata.counts.values()) == 12
    assert all(0.0 < strata.mean_alpha[n] < 1.0 for n in KIND_NAMES if strata.counts[n]), strata
    assert "alpha strata: clean" in capsys.readouterr().out


def test_chunk_objective_freezes_the_trained_head_and_saves_g_alpha(tmp_path):
    """GraphTrainer under CHUNK freezes Stage 1's head, optimises alpha alone, and writes g_alpha.npz.

    Three ways this silently produces a useless artifact: the retriever is randomly initialised
    rather than loaded, so alpha fits a posterior nobody deploys; the optimiser still holds
    retriever parameters; or the run rewrites the head it froze. All three look like success.
    """
    from onf.blend.alpha import ALPHA_NPZ, AlphaNet
    from onf.blend.ee_track import ActionScale
    from onf.blend.target import ChunkTargetBuilder
    from onf.graph.train.loop import GraphTrainer
    from onf.graph.train.loss import CHUNK

    nodes, edges = _toy_graph(n_tasks=2, demos_per_task=4, length=30)
    cfg = GraphConfig(hidden=16, layers=2, device="cpu")
    head = _small_net(nodes)
    head_bytes = Path(
        save_checkpoint(head, nodes, edges, tmp_path / schema.HEAD_NPZ)
    ).read_bytes()

    trainer = GraphTrainer(tmp_path, TrainGraphSpec(cfg=cfg, objective=CHUNK, device="cpu"))
    trainer.nodes, trainer.edges = nodes, edges
    # Only chunk_len is read below; the builder's own inputs belong to the recorded corpus.
    trainer.chunk = ChunkTargetBuilder(
        nodes=nodes, scale=ActionScale(pos=0.01, rot=0.1, r2_pos=1.0, r2_rot=1.0),
        policy=None, kinematics=None, chunk_len=5,
    )
    trainer._build_model()

    assert not any(p.requires_grad for p in trainer.net.parameters())
    for k, v in head.state_dict().items():
        assert torch.allclose(v, trainer.net.state_dict()[k]), f"{k}: not the trained head"
    owned = {id(p) for group in trainer.opt.param_groups for p in group["params"]}
    assert owned == {id(p) for p in trainer.alpha.parameters()}

    trainer.save()
    assert (tmp_path / schema.HEAD_NPZ).read_bytes() == head_bytes, "the frozen head was rewritten"
    loaded = AlphaNet.load(tmp_path / ALPHA_NPZ)
    assert (loaded.embed_dim, loaded.chunk_len) == (cfg.hidden, 5)
    for k, v in trainer.alpha.state_dict().items():
        assert torch.equal(v, loaded.state_dict()[k]), k


def test_a_policy_corruptor_actually_reaches_the_cached_chunks():
    """The wiring, not the sampler: PolicyCorruptor -> QuerySet -> ChunkTargets.a_policy.

    Every piece of this can be correct while the corruption never reaches the array the loss reads,
    and the symptom would be an artifact that trains cleanly and contains no change at all.
    """
    from onf.blend.ee_track import ActionScale
    from onf.blend.kinematics import PandaKinematics
    from onf.blend.target import ChunkTargetBuilder, PolicyActionCache
    from onf.graph.train.corrupt import P_CLEAN, PolicyCorruptor, RetrievalCorruptor
    from onf.graph.train.loop import QuerySet

    # dim=7: ChunkTargetBuilder reads nodes.ee_segment, which forward-kinematicks a Panda.
    nodes, _edges = _toy_graph(n_tasks=3, demos_per_task=3, length=40, dim=7)
    n_raw = int(nodes.n_raw)
    rng = np.random.RandomState(0)
    cache = PolicyActionCache(
        a_pi_raw=np.tile(np.arange(n_raw, dtype=np.float32)[:, None], (1, 7)),
        raw_ptr=np.asarray(nodes.raw_ptr, dtype=np.int64),
    )

    builder = ChunkTargetBuilder(
        nodes=nodes, scale=ActionScale(pos=0.0116, rot=0.1092, r2_pos=1.0, r2_rot=1.0),
        policy=cache, kinematics=PandaKinematics(), chunk_len=4,
    )
    queries = make_queries(nodes, rng, 256, np.arange(nodes.n_demos), QuerySpec(window=8))
    corruptor = PolicyCorruptor(nodes, RetrievalCorruptor(nodes))

    clean = QuerySet.from_queries(queries, torch.device("cpu"), builder)
    dirty = QuerySet.from_queries(
        queries, torch.device("cpu"), builder, rng=np.random.default_rng(0),
        policy_corruptor=corruptor,
    )
    kind = dirty.policy_corrupt.kind
    assert (kind != P_CLEAN).sum() > 40, "the sampler produced nothing to check"

    moved = ~np.all(clean.targets.a_policy == dirty.targets.a_policy, axis=(1, 2))
    # Every corrupted row moved, and no clean row did. A GAIN of ~1.0 is possible but vanishingly
    # unlikely over a log-uniform draw across two decades.
    assert np.array_equal(moved, kind != P_CLEAN), (moved.sum(), int((kind != P_CLEAN).sum()))
    # The TARGET is untouched: that is what makes a corrupted-policy row push alpha UP.
    assert np.array_equal(clean.targets.a_corr, dirty.targets.a_corr)