onf/tests/test_track.py

"""Tests for onf.graph.run.track — the sequential Bayesian belief filter over the demonstration
graph (the "WHERE: belief filter" section of onf.graph.core.schema) — plus
onf.graph.core.nodes.NodeTable.node_at_raw, the canonical raw-frame lookup TransitionKernel
is built from.

Mostly data-coupling-free (small, synthetic, in-memory graphs — mirrors tests/test_retrieve.py's
fixture pattern), except for a handful of tests gated on the real long suite artifacts at
outputs/long/frozen_learn1/artifacts (V=28476, 500 demos, 767k edges), which validate that the
vectorised node_at_raw/TransitionKernel.build construction actually holds at real scale and
report the per-check cost — skipped cleanly (tests/test_field_parity.py's skip_no_field
pattern) when that directory is absent.
"""

from __future__ import annotations

import math
import time
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.run.retrieve import GraphRetriever
from onf.graph.core.geometry import graph_hash
from onf.graph.run.track import (
    TASK_PRIOR_EPS,
    TASK_PRIOR_MIN_CHECKS,
    GraphTracker,
    TaskPosterior,
    TrackParams,
    TransitionKernel,
)
from onf.graph.run.params import TrackerCalibration
from onf.graph.build.tracker_fit import build_likelihood_sequences, evaluate_kernel_log_evidence, fit_transition_kernel

# ======================================================================================================
# real-artifact skip marker (tests/test_field_parity.py's skip_no_field pattern)
# ======================================================================================================
REAL_GRAPH_DIR = Path(__file__).resolve().parents[1] / "outputs" / "long" / "frozen_learn1" / "artifacts"
HAS_REAL_GRAPH = (REAL_GRAPH_DIR / schema.NODES_NPZ).exists()
skip_no_real_graph = pytest.mark.skipif(not HAS_REAL_GRAPH, reason=f"no real graph artifacts at {REAL_GRAPH_DIR}")


# ======================================================================================================
# synthetic graph fixture (mirrors tests/test_retrieve.py's _make_demo/_toy_nodes_edges/_toy_net)
# ======================================================================================================
DIM = 4
N_FRAMES = 31
COARSEN = 3
HIDDEN = 16


def _make_demo(task_id: int, seed: int, n: int = N_FRAMES, dim: int = DIM) -> dict:
    rng = np.random.RandomState(seed)
    start = 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]
    grip = np.zeros(n, dtype=np.uint8)
    grip[n // 2 :] = 1
    stage = grip.astype(np.int16)
    return {"q": q, "qdot": qdot, "grip": grip, "stage": stage, "task_id": task_id}


def _toy_nodes_edges(n_demos: int = 6, dim: int = DIM) -> tuple[NodeTable, EdgeSet]:
    demos = [_make_demo(task_id=i % 2, seed=100 + i, dim=dim) for i in range(n_demos)]
    nodes = NodeTable.from_demos(demos, task_names=["taskA", "taskB"], coarsen=COARSEN)
    edges = EdgeSet.build(nodes, k_sib=4, k_align=2, device="cpu")
    return nodes, edges


def _toy_net(nodes: NodeTable, hidden: int = HIDDEN, layers: int = 2, seed: int = 0) -> GraphRetrieverNet:
    torch.manual_seed(seed)
    net = GraphRetrieverNet(dim=nodes.dim, hidden=hidden, layers=layers, n_rel=schema.N_RELATIONS, agg="sum")
    mean, std = nodes.feature_stats()
    net.set_stats(mean, std)
    net.eval()
    return net


def _declare_ee_base(retriever: GraphRetriever) -> GraphRetriever:
    """Stand in for NodeTable.ee_base on a corpus that is not a 7-DoF Panda.

    ee_base is normally forward-kinematicked from q_raw, which needs a real robot; these demos are
    straight lines in a D=5 joint space. Only the readout path that GATHERS an end-effector segment
    is exercised here, never its values, so the demos' own declared ee is the right stand-in."""
    retriever.nodes.ee_base = retriever.nodes.ee_raw
    return retriever


def _toy_retriever(n_demos: int = 6, seed: int = 0) -> GraphRetriever:
    nodes, edges = _toy_nodes_edges(n_demos=n_demos)
    net = _toy_net(nodes, seed=seed)
    # readout="argmax" -- every caller of this helper feeds the retriever into a GraphTracker, whose
    # internal per-node finalize (onf.graph.run.track._finalize) only supports "argmax" (it hand-rolls the
    # top-M/softmax aggregation itself rather than delegating to onf.graph.run.readout's registered arms,
    # which are built for a raw network posterior, not a filtered belief b_t -- see that function's own
    # docstring). GraphConfig's own default is "euc_raw" (the entry-side arm), which would make
    # GraphTracker.step() raise here.
    return _declare_ee_base(
        GraphRetriever(net, nodes, edges, cfg=GraphConfig(device="cpu", readout="argmax"))
    )


# ======================================================================================================
# 1. NodeTable.node_at_raw
# ======================================================================================================
def _old_node_at_raw(nodes: NodeTable, owner: int, raw_idx: int) -> int:
    """The per-demo reference implementation this method replaces (verbatim from
    onf.graph.train.data._node_at_raw) — used here only to cross-check the vectorised implementation."""
    node_ids = np.asarray(nodes.demo_nodes(int(owner)))
    t_raws = np.asarray(nodes.t_raw)[node_ids]
    pos = int(np.searchsorted(t_raws, int(raw_idx), side="right")) - 1
    pos = min(max(pos, 0), len(node_ids) - 1)
    return int(node_ids[pos])


def test_node_at_raw_scalar_matches_per_demo_reference():
    nodes, _ = _toy_nodes_edges()
    rng = np.random.RandomState(0)
    for _ in range(200):
        j = int(rng.randint(nodes.n_demos))
        demo_start, demo_end = int(nodes.raw_ptr[j]), int(nodes.raw_ptr[j + 1])
        raw_idx = int(rng.randint(demo_start - 3, demo_end + 10))
        assert nodes.node_at_raw(j, raw_idx) == _old_node_at_raw(nodes, j, raw_idx)


def test_node_at_raw_vectorised_matches_scalar_and_is_identity_at_zero():
    nodes, _ = _toy_nodes_edges()
    owner = np.asarray(nodes.owner, dtype=np.int64)
    t_raw = np.asarray(nodes.t_raw, dtype=np.int64)

    succ0 = nodes.node_at_raw_many(owner, t_raw + 0)
    assert np.array_equal(succ0, np.arange(len(nodes)))     # a=0 is the identity map

    succ4 = nodes.node_at_raw_many(owner, t_raw + 4)
    expected = np.array([_old_node_at_raw(nodes, int(o), int(t) + 4) for o, t in zip(owner, t_raw)])
    assert np.array_equal(succ4, expected)


def test_node_at_raw_returns_python_int_for_scalar_input():
    nodes, _ = _toy_nodes_edges()
    v = nodes.node_at_raw(0, 3)
    assert isinstance(v, int)


# ======================================================================================================
# 2. TransitionKernel construction invariants
# ======================================================================================================
def test_succ_a_never_crosses_demo_boundary_and_is_monotone():
    nodes, edges = _toy_nodes_edges()
    kernel = TransitionKernel.build(nodes, edges, device="cpu")
    owner = np.asarray(nodes.owner)
    t_raw = np.asarray(nodes.t_raw)

    for i, a in enumerate(kernel.advances):
        succ = kernel.succ[i].cpu().numpy()
        assert np.array_equal(owner[succ], owner), f"advance={a} crosses a demo boundary"
        # monotone in t_raw: within any (owner) block, t_raw[succ[u]] is non-decreasing in t_raw[u]
        # (owner is itself non-decreasing/contiguous, so adjacent-index comparison suffices)
        same_owner_adjacent = owner[:-1] == owner[1:]
        t_succ = t_raw[succ]
        assert np.all(t_succ[1:][same_owner_adjacent] >= t_succ[:-1][same_owner_adjacent]), (
            f"advance={a} is not monotone in t_raw within a demo"
        )


def test_kernel_rows_sum_to_one():
    """b @ P sums to 1 for every one-hot b -- the row-stochastic property, checked directly rather than
    materialising a dense [V,V] matrix (see TransitionKernel.push's mass-conservation docstring)."""
    nodes, edges = _toy_nodes_edges()
    kernel = TransitionKernel.build(nodes, edges, beta=0.4, device="cpu")
    v = len(nodes)
    for u in [0, 1, v // 2, v - 1]:
        b = torch.zeros(v, dtype=torch.float64)
        b[u] = 1.0
        pred = kernel.push(b)
        assert pred.sum().item() == pytest.approx(1.0, abs=1e-9), f"row {u} does not sum to 1"
        assert bool((pred >= 0).all())


def test_kernel_push_conserves_mass_for_arbitrary_belief():
    nodes, edges = _toy_nodes_edges()
    kernel = TransitionKernel.build(nodes, edges, beta=0.4, device="cpu")
    rng = np.random.RandomState(0)
    raw = rng.rand(len(nodes))
    b = torch.as_tensor(raw / raw.sum(), dtype=torch.float64)
    pred = kernel.push(b)
    assert pred.sum().item() == pytest.approx(1.0, abs=1e-9)


def test_kernel_with_params_reuses_precomputed_structure():
    nodes, edges = _toy_nodes_edges()
    k1 = TransitionKernel.build(nodes, edges, device="cpu")
    k2 = k1.with_params(beta=0.9)
    assert k2.succ is k1.succ                    # same tensor object, not a copy
    assert k2.sib_src is k1.sib_src
    assert k2.beta == 0.9 and k1.beta != 0.9


def test_kernel_advance_set_and_pi_length_match_schema():
    nodes, edges = _toy_nodes_edges()
    kernel = TransitionKernel.build(nodes, edges, device="cpu")
    assert kernel.advances == schema.TRACK_ADVANCE_SET
    assert kernel.pi.shape[0] == len(schema.TRACK_ADVANCE_SET)
    assert kernel.pi.sum().item() == pytest.approx(1.0, abs=1e-9)


# ======================================================================================================
# 2b. E.1 fused kernel -- scatter-vs-fused equality, lazy build, mass conservation
# ======================================================================================================
def test_fused_push_matches_scatter_push_on_a_random_belief():
    """The equality test the plan's E.1 asks for: fuse=True's one torch.sparse.mm must reproduce
    fuse=False's ~10-scatter reference to 1e-12 on a random belief -- same pi/beta, only the
    implementation differs."""
    nodes, edges = _toy_nodes_edges()
    rng = np.random.RandomState(0)
    raw = rng.rand(len(nodes))
    b = torch.as_tensor(raw / raw.sum(), dtype=torch.float64)

    k_scatter = TransitionKernel.build(nodes, edges, beta=0.4, device="cpu", fuse=False)
    k_fused = TransitionKernel.build(nodes, edges, beta=0.4, device="cpu", fuse=True)
    pred_scatter = k_scatter.push(b)
    pred_fused = k_fused.push(b)
    np.testing.assert_allclose(pred_fused.numpy(), pred_scatter.numpy(), atol=1e-12, rtol=0)


def test_fused_build_is_lazy_and_cached():
    """fuse=True must not build the sparse matrix at build()/with_params() time -- only on the FIRST
    push() call after -- and with_params() must never carry a stale cache forward (see the class
    docstring's "FUSION" section for why fit_transition_kernel's grid search depends on this)."""
    nodes, edges = _toy_nodes_edges()
    kernel = TransitionKernel.build(nodes, edges, device="cpu", fuse=True)
    assert kernel._fused_pt is None                      # not built at construction

    k2 = kernel.with_params(beta=0.5)
    assert k2._fused_pt is None                          # with_params never inherits a stale cache
    assert k2.fuse is True                                # ... but DOES inherit the flag

    b = torch.full((len(nodes),), 1.0 / len(nodes), dtype=torch.float64)
    k2.push(b)
    assert k2._fused_pt is not None                       # built (and cached) by the first push()
    cached = k2._fused_pt
    k2.push(b)
    assert k2._fused_pt is cached                         # second push reuses the SAME cached tensor


def test_fused_kernel_rows_sum_to_one():
    """Mass conservation on the fused path -- the exact twin of test_kernel_rows_sum_to_one, fuse=True."""
    nodes, edges = _toy_nodes_edges()
    kernel = TransitionKernel.build(nodes, edges, beta=0.4, device="cpu", fuse=True)
    v = len(nodes)
    for u in [0, 1, v // 2, v - 1]:
        b = torch.zeros(v, dtype=torch.float64)
        b[u] = 1.0
        pred = kernel.push(b)
        assert pred.sum().item() == pytest.approx(1.0, abs=1e-9), f"fused row {u} does not sum to 1"
        assert bool((pred >= 0).all())


def test_fused_push_ignores_precomputed_b_mid():
    """The fused path already encodes the sibling mix in P -- passing a (deliberately WRONG) b_mid must
    change nothing, unlike the scatter path where it is consumed."""
    nodes, edges = _toy_nodes_edges()
    v = len(nodes)
    b = torch.full((v,), 1.0 / v, dtype=torch.float64)
    bogus_b_mid = torch.zeros(v, dtype=torch.float64)
    bogus_b_mid[0] = 1.0

    kernel = TransitionKernel.build(nodes, edges, beta=0.4, device="cpu", fuse=True)
    pred_default = kernel.push(b)
    pred_with_bogus = kernel.push(b, b_mid=bogus_b_mid)
    np.testing.assert_array_equal(pred_default.numpy(), pred_with_bogus.numpy())


# ======================================================================================================
# 4. backward-jump suppression
# ======================================================================================================
def test_backward_jump_is_suppressed():
    """Small synthetic two-branch aliasing graph: after the belief has advanced along demo 0, a query
    that matches an early-phase node on a DIFFERENT branch (demo 1) EQUALLY well (same L_t mass) as the
    dynamics-consistent continuation must NOT capture the belief -- see
    GraphTracker._update_belief_core's docstring for the mechanism (pred has EXACTLY zero mass on a node
    no forward advance can reach, so only the small leak term feeds it, while the dynamics-consistent
    node gets the full (1-leak)*pred term too)."""
    nodes, edges = _toy_nodes_edges()
    kernel = TransitionKernel.build(nodes, edges, device="cpu")     # default uniform pi, beta=0.3
    retriever = _toy_retriever()
    tracker = GraphTracker(retriever, kernel, calib=TrackerCalibration(leak=0.1))

    demo0 = nodes.demo_nodes(0)
    demo1 = nodes.demo_nodes(1)
    u_prev = int(demo0[3])            # belief currently sits mid-early in demo 0
    e_node = int(demo1[0])            # the aliased early-phase node -- a DIFFERENT branch entirely

    b_prev = torch.zeros(len(nodes), dtype=torch.float64)
    b_prev[u_prev] = 1.0
    pred = kernel.push(b_prev)
    assert pred[e_node].item() == 0.0, "pred must never route mass backward onto a different branch"
    c_node = int(torch.argmax(pred).item())

    L_t = torch.zeros(len(nodes), dtype=torch.float64)
    L_t[e_node] = 0.5
    L_t[c_node] = 0.5                 # the observation, on its own, cannot tell the two apart

    b_t, _, _, _ = tracker._update_belief_core(b_prev, L_t, len(nodes))
    assert b_t[c_node].item() > b_t[e_node].item(), (
        "the aliased early-phase node captured the belief despite the dynamics giving it zero support"
    )
    assert b_t[e_node].item() < 0.5, "L_t alone gave e_node half the mass; the filter must reduce that"


# ======================================================================================================
# 8. g_track.npz — graph_hash refusal
# ======================================================================================================
def test_g_track_npz_round_trips(tmp_path):
    nodes, edges = _toy_nodes_edges()
    params = TrackParams(
        pi=np.array([0.5, 0.1, 0.1, 0.1, 0.1, 0.05, 0.025, 0.025]), beta=0.3, leak=0.1,
    )
    out = params.save(tmp_path / "graph", nodes, edges)
    loaded = TrackParams.load(out, nodes, edges)
    np.testing.assert_allclose(loaded.pi, params.pi)
    assert loaded.beta == pytest.approx(params.beta)
    assert loaded.leak == pytest.approx(params.leak)
    assert math.isnan(loaded.delta_star)                              # reserved-slot placeholder
    assert loaded.s_ref_per_joint.shape == (nodes.dim,)
    assert np.all(np.isnan(loaded.s_ref_per_joint))
    # Basin slots unset here, so they round-trip as the documented NaN placeholders.
    assert loaded.basin_r.shape == (nodes.n_tasks, schema.NBINS_ALIGN)
    assert np.all(np.isnan(loaded.basin_r))
    assert loaded.basin_h.shape == (nodes.n_tasks, schema.NBINS_ALIGN)
    assert np.all(np.isnan(loaded.basin_h))


def test_g_track_npz_round_trips_with_basin_fields(tmp_path):
    """The basin slots are POPULATED this time -- pin that real geometry survives the round trip
    bit-for-bit, not just the empty-placeholder path the previous test exercises."""
    nodes, edges = _toy_nodes_edges()
    n_tasks, nbins = nodes.n_tasks, schema.NBINS_ALIGN
    rng = np.random.RandomState(0)
    basin_r = rng.uniform(0.5, 1.5, size=(n_tasks, nbins))
    basin_h = rng.uniform(0.1, 0.3, size=(n_tasks, nbins))
    params = TrackParams(pi=np.full(8, 1 / 8), beta=0.3, leak=0.1, basin_r=basin_r, basin_h=basin_h)
    out = params.save(tmp_path / "graph", nodes, edges)
    loaded = TrackParams.load(out, nodes, edges)
    np.testing.assert_allclose(loaded.basin_r, basin_r)
    np.testing.assert_allclose(loaded.basin_h, basin_h)


def test_g_track_npz_load_ignores_retired_keys(tmp_path):
    """Files written before the null banks were retired still carry those arrays; load must ignore
    the extra keys rather than fail."""
    nodes, edges = _toy_nodes_edges()
    out = TrackParams(pi=np.full(8, 1 / 8), beta=0.3, leak=0.1).save(tmp_path / "graph", nodes, edges)
    payload = dict(np.load(out))
    payload["null_bank_evidence"] = np.array([-1.0, 0.0, 1.0])
    payload["null_bank_progress"] = np.array([0.1, 0.5, 0.9])
    payload["null_bank_basin"] = np.array([-6.0, -3.0, -1.0])
    np.savez(out, **payload)

    loaded = TrackParams.load(out, nodes, edges)
    assert loaded.beta == pytest.approx(0.3)
    assert not hasattr(loaded, "null_bank_evidence")


def test_g_track_npz_graph_hash_mismatch_raises(tmp_path):
    nodes, edges = _toy_nodes_edges()
    params = TrackParams(pi=np.full(8, 1 / 8), beta=0.3, leak=0.1)
    out_dir = tmp_path / "graph"
    out_dir.mkdir()
    payload = {
        "pi": params.pi, "beta": np.array(params.beta), "leak": np.array(params.leak),
        "delta_star": np.array(float("nan")), "s_ref_per_joint": np.full(nodes.dim, np.nan),
        "graph_hash": np.array("not-the-right-hash"),
    }
    np.savez(out_dir / schema.TRACK_NPZ, **payload)

    with pytest.raises(ValueError, match="graph_hash mismatch"):
        TrackParams.load(out_dir, nodes, edges)


def test_g_track_npz_missing_graph_hash_raises(tmp_path):
    nodes, edges = _toy_nodes_edges()
    out_dir = tmp_path / "graph"
    out_dir.mkdir()
    payload = {
        "pi": np.full(8, 1 / 8), "beta": np.array(0.3), "leak": np.array(0.1),
        "delta_star": np.array(float("nan")), "s_ref_per_joint": np.full(nodes.dim, np.nan),
    }
    np.savez(out_dir / schema.TRACK_NPZ, **payload)          # no graph_hash key at all

    with pytest.raises(ValueError, match="graph_hash mismatch"):
        TrackParams.load(out_dir, nodes, edges)


def test_graph_tracker_load_refuses_stale_track_params(tmp_path):
    """End-to-end: GraphTracker.load must propagate the same refusal when a saved graph directory's
    g_track.npz does not match its own g_nodes/g_edges."""
    nodes, edges = _toy_nodes_edges()
    net = _toy_net(nodes)
    d = tmp_path / "graph"
    d.mkdir()
    nodes.save(d)
    edges.save(d / schema.EDGES_NPZ)
    head = dict(net.state_npz_dict())
    head["graph_hash"] = np.array(graph_hash(nodes, edges))
    np.savez(d / schema.HEAD_NPZ, **head)

    bad_payload = {
        "pi": np.full(8, 1 / 8), "beta": np.array(0.3), "leak": np.array(0.1),
        "delta_star": np.array(float("nan")), "s_ref_per_joint": np.full(nodes.dim, np.nan),
        "graph_hash": np.array("wrong"),
    }
    np.savez(d / schema.TRACK_NPZ, **bad_payload)

    with pytest.raises(ValueError, match="graph_hash mismatch"):
        GraphTracker.load(d, cfg=GraphConfig(device="cpu"))


def test_graph_tracker_load_without_track_npz_falls_back_to_unfit_defaults(tmp_path):
    nodes, edges = _toy_nodes_edges()
    net = _toy_net(nodes)
    d = tmp_path / "graph"
    d.mkdir()
    nodes.save(d)
    edges.save(d / schema.EDGES_NPZ)
    head = dict(net.state_npz_dict())
    head["graph_hash"] = np.array(graph_hash(nodes, edges))
    np.savez(d / schema.HEAD_NPZ, **head)                      # no g_track.npz written at all

    # readout="argmax" -- this test calls tracker.step() below, which needs it (see _toy_retriever's
    # own comment for why).
    tracker = GraphTracker.load(d, cfg=GraphConfig(device="cpu", readout="argmax"))
    q0 = tracker.retriever.nodes.q[0].astype(np.float64)
    q_hist = np.repeat(q0[None, :], tracker.retriever.cfg.hist, axis=0)
    tracker.reset()
    step = tracker.step(q_hist)                                  # must not raise
    assert step.as_retrieval().node >= 0


# ======================================================================================================
# 8b. head_npz threading -- default bit-parity, and GraphTracker.where_target
# ======================================================================================================
def _write_toy_graph_dir(tmp_path, nodes, edges, net, *, head_name: str = schema.HEAD_NPZ) -> Path:
    d = tmp_path / "graph"
    d.mkdir(exist_ok=True)
    nodes.save(d)
    edges.save(d / schema.EDGES_NPZ)
    head = dict(net.state_npz_dict())
    head["graph_hash"] = np.array(graph_hash(nodes, edges))
    np.savez(d / head_name, **head)
    return d


def test_head_npz_default_is_bit_identical_to_explicit_default(tmp_path):
    """schema.HEAD_NPZ (the default) must be reachable two ways -- omitting head_npz entirely, or passing
    it explicitly -- and produce a bit-identical retriever; every existing call site (which never passes
    head_npz) is therefore unaffected by this parameter's addition."""
    nodes, edges = _toy_nodes_edges()
    net = _toy_net(nodes)
    d = _write_toy_graph_dir(tmp_path, nodes, edges, net)

    r_default = GraphRetriever.load(d, cfg=GraphConfig(device="cpu"))
    r_explicit = GraphRetriever.load(d, cfg=GraphConfig(device="cpu"), head_npz=schema.HEAD_NPZ)
    assert r_default.graph_hash == r_explicit.graph_hash
    for p1, p2 in zip(r_default.net.parameters(), r_explicit.net.parameters()):
        assert torch.equal(p1, p2)


_ALT_HEAD_NPZ = "g_alt_head.npz"   # any name other than schema.HEAD_NPZ; exercises head_npz's own
                                    # name-selection logic, not a specific named artifact (there is only
                                    # one head artifact any more, schema.HEAD_NPZ -- see schema.py's note
                                    # by NODES_NPZ/EDGES_NPZ/HEAD_NPZ on the removed second artifact).


def test_head_npz_can_load_a_differently_named_head(tmp_path):
    """Passing a non-default name must load THAT file instead of g_head.npz -- head_npz selects the
    on-disk artifact by name, never by which weights happen to be loaded into a shared file."""
    nodes, edges = _toy_nodes_edges()
    net = _toy_net(nodes, seed=7)
    d = _write_toy_graph_dir(tmp_path, nodes, edges, net, head_name=_ALT_HEAD_NPZ)

    with pytest.raises(FileNotFoundError):
        GraphRetriever.load(d, cfg=GraphConfig(device="cpu"))         # g_head.npz was never written here
    r = GraphRetriever.load(d, cfg=GraphConfig(device="cpu"), head_npz=_ALT_HEAD_NPZ)
    assert r.graph_hash == graph_hash(nodes, edges)


def test_graph_tracker_load_threads_head_npz(tmp_path):
    nodes, edges = _toy_nodes_edges()
    net = _toy_net(nodes)
    d = _write_toy_graph_dir(tmp_path, nodes, edges, net, head_name=_ALT_HEAD_NPZ)
    tracker = GraphTracker.load(d, cfg=GraphConfig(device="cpu"), head_npz=_ALT_HEAD_NPZ)
    assert tracker.retriever.graph_hash == graph_hash(nodes, edges)


# ======================================================================================================
# 8c. GraphTracker.where_target -- belief top-K -> ReadoutContext -> the "basin" arm
# ======================================================================================================
def test_where_target_requires_a_prior_step():
    retriever = _toy_retriever()
    kernel = TransitionKernel.build(retriever.nodes, retriever.edges, device="cpu")
    tracker = GraphTracker(retriever, kernel)
    tracker.reset()
    with pytest.raises(RuntimeError, match="no step"):
        tracker.where_target(retriever.nodes.q[0].astype(np.float64))


def test_where_target_returns_no_op_inside_basin():
    """Exercises the readout.py interface this method is written against (ReadoutContext.basin_r/
    basin_h/task_id, the 'basin' arm)."""
    retriever = _toy_retriever()
    kernel = TransitionKernel.build(retriever.nodes, retriever.edges, device="cpu")
    n_tasks = retriever.nodes.n_tasks
    basin_r = np.full((n_tasks, schema.NBINS_ALIGN), 10.0)     # generous radius -> everything is "inside"
    basin_h = np.full((n_tasks, schema.NBINS_ALIGN), 1.0)
    tracker = GraphTracker(retriever, kernel, calib=TrackerCalibration(basin_r=basin_r, basin_h=basin_h))

    q0 = retriever.nodes.q[5].astype(np.float64)
    q_hist = np.repeat(q0[None, :], retriever.cfg.hist, axis=0)
    tracker.reset()
    tracker.step(q_hist)
    result = tracker.where_target(q0)
    assert result.info.get("no_op") is True


def test_where_target_aims_at_the_pushed_belief_so_the_target_lies_AHEAD_along_the_strand():
    """Pins the v3 WHERE change. where_target used to read the readout off the FILTERED belief b_t
    ("where am I now"), which makes the correction purely perpendicular to the demo -- it restores a
    pose while discarding everything the graph knows about direction of travel, so a merge could put the
    arm back on the tube still heading the wrong way. It now reads off kernel.push(b_t) ("where does
    the demo flow say I will be next check"), which is the filter's OWN predict step -- no new fitted
    quantity. The observable consequence, pinned here: the target's depth (the p-weighted phase of the
    support the readout aggregates) must be strictly GREATER than the depth the filtered belief would
    have produced, i.e. the aim point moved forward along the strand rather than sideways onto it.
    """
    from onf.graph.run import readout as gr

    retriever = _toy_retriever()
    kernel = TransitionKernel.build(retriever.nodes, retriever.edges, device="cpu")
    n_tasks = retriever.nodes.n_tasks
    # tight radius -> never a no_op, so a real target (and a real depth) is produced
    basin_r = np.full((n_tasks, schema.NBINS_ALIGN), 1e-3)
    basin_h = np.zeros((n_tasks, schema.NBINS_ALIGN))
    tracker = GraphTracker(retriever, kernel, calib=TrackerCalibration(basin_r=basin_r, basin_h=basin_h))

    # start mid-demo, so there is strand left ahead to advance INTO (phase near 1.0 would clip)
    v_mid = int(np.argmin(np.abs(np.asarray(retriever.nodes.phase) - 0.3)))
    q0 = retriever.nodes.q[v_mid].astype(np.float64)
    q_hist = np.repeat(q0[None, :], retriever.cfg.hist, axis=0)
    tracker.reset()
    tracker.step(q_hist)

    pushed = tracker.where_target(q0)

    # the pre-change behaviour, reconstructed inline from the SAME belief: read the readout off b_t
    b = tracker._b
    k = min(schema.TRACK_BELIEF_TOPK, int(b.shape[0]))
    top_v, top_idx = torch.topk(b, k, largest=True, sorted=True)
    p_filtered = np.zeros(len(retriever.nodes), dtype=np.float64)
    p_filtered[top_idx.cpu().numpy()] = top_v.cpu().numpy()
    ctx = gr.ReadoutContext(
        nodes=retriever.nodes, topm=retriever.cfg.topm, q_now=q0,
        basin=gr.BasinGeometry(r=basin_r, h=basin_h), task_id=int(retriever.nodes.task_id[v_mid]),
    )
    filtered = gr.make_readout("basin")(p_filtered, ctx)

    assert pushed.depth > filtered.depth, (
        f"pushed-belief target must aim further along the strand: depth {pushed.depth:.4f} "
        f"vs filtered {filtered.depth:.4f}"
    )


# ======================================================================================================
# 9. MLE fit improves held-out predictive log-evidence over a uniform-pi baseline
# ======================================================================================================
def test_mle_fit_improves_held_out_log_evidence():
    retriever = _toy_retriever(n_demos=6)
    held_out = [0, 2, 4]
    result = fit_transition_kernel(retriever, held_out, n_sweeps=2)
    assert result.log_evidence >= result.log_evidence_uniform_pi - 1e-9
    assert np.isfinite(result.log_evidence)
    assert result.pi.shape == (len(schema.TRACK_ADVANCE_SET),)
    assert result.pi.sum() == pytest.approx(1.0, abs=1e-6)
    assert 0.0 <= result.beta <= 1.0
    assert 0.0 <= result.leak <= 1.0
    assert result.n_checks > 0


def test_fit_transition_kernel_raises_on_no_usable_held_out_demos():
    retriever = _toy_retriever(n_demos=6)
    with pytest.raises(ValueError, match="no usable fitting sequence"):
        fit_transition_kernel(retriever, held_out_owners=[], check_every=8)


# ======================================================================================================
# 9b. build_likelihood_sequences -- clean is not the only regime that can be fit on (2026-07-30 addition)
# ======================================================================================================
# Longer toy demos than the DIM=4/N_FRAMES=31 fixture above: a real drift/cross-strand rollout needs
# room for several checks on top of DriftConfig's own window/check-spacing floor.
_LONG_N_FRAMES = 80


def _toy_retriever_for_sequences(n_demos: int = 8, seed: int = 0) -> GraphRetriever:
    nodes, edges = _toy_nodes_edges(n_demos=n_demos)
    # rebuild with longer demos (the module-level _make_demo/_toy_nodes_edges default to N_FRAMES=31,
    # too short for a multi-check drift/cross-strand rollout under the default check_every=8)
    demos = [_make_demo(task_id=i % 2, seed=100 + i, n=_LONG_N_FRAMES) for i in range(n_demos)]
    nodes = NodeTable.from_demos(demos, task_names=["taskA", "taskB"], coarsen=COARSEN)
    edges = EdgeSet.build(nodes, k_sib=4, k_align=2, device="cpu")
    net = _toy_net(nodes, seed=seed)
    # readout="argmax" -- see _toy_retriever's comment; every caller here feeds into GraphTracker too.
    return GraphRetriever(net, nodes, edges, cfg=GraphConfig(device="cpu", readout="argmax"))


def _assert_valid_L_sequences(sequences, v):
    assert len(sequences) > 0
    for L in sequences:
        assert L.ndim == 2 and L.shape[0] >= 2 and L.shape[1] == v
        row_sums = L.sum(dim=1)
        # every row is either a valid distribution (sums to 1) or all-zero (nothing reachable under
        # this window -- should not happen for these toy queries, but a hard 1.0 assertion here would
        # make the test brittle to unrelated toy-graph changes; check the common case strictly instead.
        assert torch.all((row_sums > 0.999) & (row_sums < 1.001))


def test_build_likelihood_sequences_clean():
    retriever = _toy_retriever_for_sequences()
    seqs = build_likelihood_sequences(retriever, [0, 2, 4], kind="clean", check_every=8)
    _assert_valid_L_sequences(seqs, len(retriever.nodes))
    assert len(seqs) <= 3          # at most one sequence per requested owner


def test_build_likelihood_sequences_drift():
    retriever = _toy_retriever_for_sequences()
    seqs = build_likelihood_sequences(
        retriever, [0, 2, 4, 6], kind="drift", n_sequences=8, check_every=8, seed=1,
    )
    _assert_valid_L_sequences(seqs, len(retriever.nodes))


def test_build_likelihood_sequences_cross_strand():
    retriever = _toy_retriever_for_sequences()
    seqs = build_likelihood_sequences(
        retriever, [0, 2, 4, 6], kind="cross_strand", n_sequences=8, check_every=8, seed=2,
    )
    _assert_valid_L_sequences(seqs, len(retriever.nodes))


def test_build_likelihood_sequences_rejects_unknown_kind():
    retriever = _toy_retriever_for_sequences()
    with pytest.raises(ValueError, match="kind must be one of"):
        build_likelihood_sequences(retriever, [0], kind="not_a_real_kind")


def test_build_likelihood_sequences_empty_owners_returns_empty_list():
    retriever = _toy_retriever_for_sequences()
    assert build_likelihood_sequences(retriever, [], kind="clean") == []
    assert build_likelihood_sequences(retriever, [], kind="drift") == []


def test_fit_transition_kernel_accepts_sequences_directly():
    """The 2026-07-30 fix this whole section exists for: fitting is no longer hard-wired to
    clean-demo-only sequences -- a caller can hand it ANY [T,V] sequence set, e.g. drift or
    cross-strand, and the fit optimises against exactly that set."""
    retriever = _toy_retriever_for_sequences()
    drift_seqs = build_likelihood_sequences(
        retriever, [0, 2, 4, 6], kind="drift", n_sequences=6, check_every=8, seed=3,
    )
    result = fit_transition_kernel(retriever, sequences=drift_seqs, n_sweeps=1)
    assert result.log_evidence >= result.log_evidence_uniform_pi - 1e-9
    assert result.n_checks == sum(L.shape[0] - 1 for L in drift_seqs)


def test_fit_transition_kernel_combines_sequences_and_held_out_owners():
    retriever = _toy_retriever_for_sequences()
    drift_seqs = build_likelihood_sequences(
        retriever, [0, 2], kind="drift", n_sequences=3, check_every=8, seed=4,
    )
    result = fit_transition_kernel(retriever, held_out_owners=[4, 6], sequences=drift_seqs, n_sweeps=1)
    clean_seqs = build_likelihood_sequences(retriever, [4, 6], kind="clean", check_every=8)
    assert result.n_checks == sum(L.shape[0] - 1 for L in drift_seqs) + sum(L.shape[0] - 1 for L in clean_seqs)


def test_evaluate_kernel_log_evidence_matches_fit_internal_computation():
    """evaluate_kernel_log_evidence is the SAME quantity fit_transition_kernel's coordinate ascent
    scores internally -- pin that a standalone call against a FIXED kernel reproduces exactly what the
    fit found at its own optimum (beta=0 as a probe point, not assuming the fit landed there)."""
    retriever = _toy_retriever_for_sequences()
    seqs = build_likelihood_sequences(
        retriever, [0, 2, 4, 6], kind="cross_strand", n_sequences=6, check_every=8, seed=5,
    )
    kernel = TransitionKernel.build(retriever.nodes, retriever.edges, beta=0.0, device="cpu")
    ev, n = evaluate_kernel_log_evidence(kernel, leak=0.1, sequences=seqs)
    assert np.isfinite(ev)
    assert n == sum(L.shape[0] - 1 for L in seqs)


# ======================================================================================================
# 10. real-artifact-gated: scale, cost, and end-to-end sanity on the actual long graph
# ======================================================================================================
@skip_no_real_graph
def test_real_graph_transition_kernel_builds_and_succ_a_is_well_formed():
    nodes = NodeTable.load(REAL_GRAPH_DIR / schema.NODES_NPZ)
    edges = EdgeSet.load(REAL_GRAPH_DIR / schema.EDGES_NPZ)
    kernel = TransitionKernel.build(nodes, edges, device="cpu")

    owner = np.asarray(nodes.owner)
    for i, a in enumerate(kernel.advances):
        succ = kernel.succ[i].cpu().numpy()
        assert np.array_equal(owner[succ], owner), f"advance={a} crosses a demo boundary on the real graph"
    assert kernel.succ.shape == (len(schema.TRACK_ADVANCE_SET), len(nodes))


@skip_no_real_graph
def test_real_graph_kernel_push_per_check_cost():
    """Measures TransitionKernel.push's per-check cost on the real V=28476 graph -- the module
    docstring's performance budget target (~1M float ops, 0.1-1% of observe()'s ~150M). Reported via
    print (captured by pytest -s / -q -rA), not asserted tightly (wall-clock is not a portability-safe
    thing to pin), but bounded generously so a real regression still fails the test."""
    nodes = NodeTable.load(REAL_GRAPH_DIR / schema.NODES_NPZ)
    edges = EdgeSet.load(REAL_GRAPH_DIR / schema.EDGES_NPZ)
    device = "cuda" if torch.cuda.is_available() else "cpu"
    kernel = TransitionKernel.build(nodes, edges, device=device)
    v = len(nodes)
    b = torch.full((v,), 1.0 / v, dtype=torch.float64, device=device)

    kernel.push(b)  # warm-up (cuda context / caching allocator)
    if device == "cuda":
        torch.cuda.synchronize()
    n_reps = 50
    t0 = time.perf_counter()
    for _ in range(n_reps):
        pred = kernel.push(b)
    if device == "cuda":
        torch.cuda.synchronize()
    dt = (time.perf_counter() - t0) / n_reps
    print(f"\n[test_real_graph_kernel_push_per_check_cost] V={v} device={device} push()={dt * 1e3:.3f} ms/call")
    assert dt < 0.5, f"push() took {dt * 1e3:.1f} ms/call on V={v} -- unexpectedly slow"
    assert pred.sum().item() == pytest.approx(1.0, abs=1e-6)


@skip_no_real_graph
def test_real_graph_kernel_rows_sum_to_one():
    nodes = NodeTable.load(REAL_GRAPH_DIR / schema.NODES_NPZ)
    edges = EdgeSet.load(REAL_GRAPH_DIR / schema.EDGES_NPZ)
    kernel = TransitionKernel.build(nodes, edges, beta=0.3, device="cpu")
    v = len(nodes)
    rng = np.random.RandomState(0)
    for u in rng.choice(v, size=20, replace=False):
        b = torch.zeros(v, dtype=torch.float64)
        b[int(u)] = 1.0
        pred = kernel.push(b)
        assert pred.sum().item() == pytest.approx(1.0, abs=1e-6)


@skip_no_real_graph
def test_real_graph_step_overhead_vs_retrieve():
    """E.2's acceptance gate: GraphTracker.step()'s SINGLE-SYNC belief-tracking machinery, fused kernel
    included, must cost <= 10% over plain retrieve() on the real long graph (V=28476) -- reported via
    print (captured by -s/-rA), not tightly asserted (wall-clock is not portability-safe to pin),
    but bounded generously so a real regression still fails. Measured on this machine's CPU (no GPU in
    this sandbox): step() actually comes in FASTER than retrieve() here, consistent with the module
    docstring's performance budget (observe()'s ~150M FLOPs dominate the belief update's ~1M).

    Loads TWO retrievers off the same on-disk graph, one per readout: onf.graph.run.track._finalize (the
    belief-filtered aggregation tracker.step() uses) only supports cfg.readout == "argmax", while
    retriever.retrieve()'s registered arms are ("euc_raw", "basin") (`onf.graph.core.schema.
    READOUT_ARMS`) -- "argmax" is not one of them any more. This is still a fair overhead comparison:
    both retrievers share the identical trained network/nodes/edges, and observe() (~150M FLOPs) is
    the cost this measures, dominating either aggregation (~1M FLOPs) by 2 orders of magnitude.
    """
    retriever = GraphRetriever.load(REAL_GRAPH_DIR, cfg=GraphConfig(device="cpu"))
    retriever_argmax = GraphRetriever.load(REAL_GRAPH_DIR, cfg=GraphConfig(device="cpu", readout="argmax"))
    nodes = retriever.nodes
    q0 = nodes.q[100].astype(np.float64)
    q_hist = np.repeat(q0[None, :], retriever.cfg.hist, axis=0)

    n_reps = 20
    retriever.retrieve(q_hist)                                  # warm-up
    t0 = time.perf_counter()
    for _ in range(n_reps):
        retriever.retrieve(q_hist)
    t_retrieve = (time.perf_counter() - t0) / n_reps

    kernel = TransitionKernel.build(nodes, retriever_argmax.edges, device="cpu", fuse=True)
    tracker = GraphTracker(retriever_argmax, kernel)
    tracker.reset()
    tracker.step(q_hist)                                        # ENTRY step (no transition applied)
    tracker.step(q_hist)                                         # a real step -- also builds+caches the
                                                                  # fused matrix, so it is warmed up too
    t0 = time.perf_counter()
    for _ in range(n_reps):
        tracker.step(q_hist)
    t_step = (time.perf_counter() - t0) / n_reps

    overhead_pct = 100.0 * (t_step - t_retrieve) / t_retrieve
    print(
        f"\n[test_real_graph_step_overhead_vs_retrieve] V={len(nodes)} retrieve()={t_retrieve * 1e3:.3f}ms "
        f"step()={t_step * 1e3:.3f}ms overhead={overhead_pct:.1f}%"
    )
    assert overhead_pct <= 25.0, f"step() overhead {overhead_pct:.1f}% exceeds the E.2 acceptance gate"


# ======================================================================================================
# 11. TaskPosterior -- the behavioural task lane (GR_TASK_PRIOR_W)
# ======================================================================================================
def _task_posterior(weight: float, task_id=(0, 0, 1, 1, 2, 2)) -> TaskPosterior:
    return TaskPosterior(torch.as_tensor(task_id, dtype=torch.int64), max(task_id) + 1, weight=weight)


def _belief(mass: list[float]) -> torch.Tensor:
    b = torch.as_tensor(mass, dtype=torch.float64)
    return b / b.sum()


def test_task_posterior_gates_and_stays_soft():
    """The marginal is a plain per-task sum of b_t, and the prior it yields is withheld on all three
    gates (off / too little evidence / no opinion) and is FINITE when it does apply -- a -inf entry
    would make it the hard mask this deliberately is not."""
    confident = _belief([0.002, 0.002, 0.0, 0.0, 0.5, 0.496])
    uniform = _belief([1.0] * 6)
    enough = TASK_PRIOR_MIN_CHECKS

    tp = _task_posterior(1.0)
    assert tp.marginal(confident).tolist() == pytest.approx([0.004, 0.0, 0.996])

    assert _task_posterior(0.0).log_prior(confident, enough) is None      # off by default
    assert tp.log_prior(None, enough) is None                             # no belief yet
    assert tp.log_prior(confident, enough - 1) is None                    # too little evidence
    assert torch.equal(tp.log_prior(uniform, enough), torch.zeros(3, dtype=torch.float64))

    prior = tp.log_prior(confident, enough)
    assert bool(torch.isfinite(prior).all())
    assert int(torch.argmax(prior).item()) == 2
    # A task with exactly zero belief mass is floored at -log(eps), not -inf.
    assert prior[1].item() == pytest.approx(math.log(TASK_PRIOR_EPS))
    # ...and the winner is worth a few nats over it, at weight 1.
    assert 3.0 <= (prior[2] - prior[1]).item() <= 10.0


def test_task_prior_is_inert_when_the_text_lane_resolves():
    """Strictly additive: a large GR_TASK_PRIOR_W must reproduce the weight-0 run bit-for-bit whenever
    the lane resolved from the instruction string, and must be withheld while the marginal is still
    undecided. The last block hands the filter a marginal that IS decided, so the equalities above
    cannot be passing merely because nothing is wired up."""
    nodes, edges = _toy_nodes_edges()
    net = _toy_net(nodes)
    kernel = TransitionKernel.build(nodes, edges, device="cpu")

    def tracker_at(weight: float) -> GraphTracker:
        cfg = GraphConfig(device="cpu", readout="argmax", task_prior_w=weight)
        return GraphTracker(GraphRetriever(net, nodes, edges, cfg=cfg), kernel)

    def run(weight: float, task: str) -> list[np.ndarray]:
        tracker = tracker_at(weight)
        rng = np.random.RandomState(7)
        q_hist = np.asarray(nodes.q[: tracker.retriever.cfg.hist], dtype=np.float64)
        out = []
        for _ in range(2 * TASK_PRIOR_MIN_CHECKS):
            out.append(tracker.step(q_hist, task=task).as_retrieval().p_node)
            q_hist = q_hist + 0.01 * rng.randn(*q_hist.shape)
        return out

    for task in ("taskA", ""):
        # "taskA" resolves the lane, so the prior is dropped; "" resolves nothing, but the untrained
        # fixture head leaves the marginal at normalised entropy ~0.98, so the gate withholds it.
        for off, on in zip(run(0.0, task), run(8.0, task)):
            np.testing.assert_array_equal(off, on)

    def one_step_on_a_committed_belief(weight: float, task: str) -> np.ndarray:
        tracker = tracker_at(weight)
        # Reaching normalised entropy < TASK_PRIOR_MAX_ENTROPY takes a TRAINED head and ~8 real checks
        # (measured on the long tapes); the fixture head never gets there. What is under test here is
        # the wiring and the lane precedence, not the head, so the belief is handed over directly.
        b = torch.as_tensor(np.asarray(nodes.task_id) == 1, dtype=torch.float64)
        tracker._b = b / b.sum()
        tracker._n_checks = TASK_PRIOR_MIN_CHECKS
        q_hist = np.asarray(nodes.q[: tracker.retriever.cfg.hist], dtype=np.float64)
        return tracker.step(q_hist, task=task).as_retrieval().p_node

    np.testing.assert_array_equal(
        one_step_on_a_committed_belief(0.0, "taskA"), one_step_on_a_committed_belief(8.0, "taskA"),
    )
    assert not np.array_equal(
        one_step_on_a_committed_belief(0.0, ""), one_step_on_a_committed_belief(8.0, ""),
    ), "a committed marginal with no text lane changed nothing -- the prior is not wired up"