onf/tests/test_readout.py

"""Tests for onf.graph.run.readout — synthetic-only, runs in .venv-core (numpy + torch, no HDF5,
no trained GNN). The retrieval distribution p(v|Q) is hand-built for every test; the module under
test only cares about aggregating a GIVEN p, never about how p was produced.

This repo ships exactly two registered arms: euc_raw (deployed t=0 ENTRY) and basin (deployed
t>0 SENTINEL); every test below exercises only these two.
"""

from __future__ import annotations

import numpy as np
import pytest
import torch

from onf.graph.core import schema
from onf.graph.core.nodes import NodeTable
from onf.graph.core.segments import ee_segment_table, expected_ee_segment
from onf.graph.run.readout import (
    BasinGeometry,
    ARMS,
    TOPK_K_MAX,
    TOPK_K_MIN,
    ReadoutContext,
    ReadoutResult,
    make_readout,
)

D = 5     # joint-space dim used throughout


# ======================================================================================================
# synthetic-context helpers
# ======================================================================================================
def _ee_from_q(q: np.ndarray) -> np.ndarray:
    """[T, 6] end-effector track built as q's first 3 columns twice, so a posterior-weighted ee_seg
    is comparable column-for-column against the q_seg carrying the same weights."""
    return np.concatenate([q[:, :3], q[:, :3]], axis=1).astype(np.float32)


def _declare_ee_base(nodes: NodeTable) -> NodeTable:
    """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
    D=5 random walks. The arms under test only ever ask whether the SAME table flows through the
    same support with the same weights, so the demos' own declared ee is the right stand-in."""
    nodes.ee_base = nodes.ee_raw
    return nodes


def _make_demo(t: int, seed: int, task_id: int = 0, offset: float = 0.0) -> dict:
    """A deterministic, smoothly-moving demo: single stage, gripper open throughout (keeps
    coarsen=1 node count exactly equal to frame count, so node indices are hand-traceable)."""
    rng = np.random.RandomState(seed)
    steps = rng.randn(t, D).astype(np.float32) * 0.01
    q = (np.cumsum(steps, axis=0) + offset).astype(np.float32)
    qdot = np.zeros_like(q)
    qdot[:-1] = q[1:] - q[:-1]
    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,
            "ee": _ee_from_q(q)}


def _flat_context(n_demos: int = 3, t: int = 12, topm: int = 32) -> tuple[NodeTable, ReadoutContext]:
    """A plain multi-demo table — enough for euc_raw/basin."""
    demos = [_make_demo(t, seed=i) for i in range(n_demos)]
    nodes = _declare_ee_base(NodeTable.from_demos(demos, task_names=["taskA"], coarsen=1))
    return nodes, ReadoutContext(nodes=nodes, topm=topm)


def _fork_context(fork: int = 4, t: int = 12, topm: int = 32) -> tuple[NodeTable, ReadoutContext, int, int]:
    """Two demos sharing an IDENTICAL trunk (t=0..fork) then diverging in joint space
    (t=fork+1..T-1) — enough geometric structure to exercise "mass spread over two distinct configs"
    without needing any latent embedding (the arms that consumed one -- hyp/euc_embed -- are
    gone; see the module docstring).

    Returns (nodes, ctx, tip0, tip1) — the two branch-tip node indices mass gets put on.
    """
    demos = [_make_demo(t, seed=0, offset=0.0), _make_demo(t, seed=1, offset=5.0)]
    nodes = _declare_ee_base(NodeTable.from_demos(demos, task_names=["taskA"], coarsen=1))
    ctx = ReadoutContext(nodes=nodes, topm=topm)
    tip0 = int(nodes.demo_ptr[0]) + t - 1
    tip1 = int(nodes.demo_ptr[1]) + t - 1
    return nodes, ctx, tip0, tip1


def _concentrated_p(v: int, n: int) -> np.ndarray:
    """Mass ~entirely on node v, with a tiny positive floor everywhere else so weighted_topk's
    k_min filler slots are well-defined (zero weight, but not literally absent probabilities)."""
    p = np.full(n, 1e-9)
    p[v] = 1.0
    return p / p.sum()


def _fixed_demo(q_row: np.ndarray, task_id: int = 0) -> dict:
    """A 2-identical-frame demo holding EXACTLY the given joint config at every frame -- lets a test
    hand-pick a candidate q-value precisely (unlike the random walk in _make_demo), while still
    giving NodeTable.from_demos enough frames (T=2, so phase spans exactly {0, 1}) to
    satisfy the table's own phase-convention invariant (NodeTable.validate requires table-wide
    phase max > 0.99 -- a single-frame demo's phase is always exactly 0 and fails that check)."""
    q = np.tile(np.asarray(q_row, dtype=np.float32)[None, :], (2, 1))
    qdot = np.zeros_like(q)
    grip = np.zeros(2, dtype=np.uint8)
    stage = np.zeros(2, dtype=np.int16)
    return {"q": q, "qdot": qdot, "grip": grip, "stage": stage, "task_id": task_id,
            "ee": _ee_from_q(q)}


def _two_node_context(
    q0: np.ndarray, q1: np.ndarray, topm: int = 32
) -> tuple[NodeTable, ReadoutContext, int, int]:
    """A graph built from two 2-frame demos (coarsen=1, so each frame is its own node): the LAST
    frame of demo 0/1 (phase == 1.0, table indices 1 and 3) carries q0/q1 exactly and is what
    every test below puts probability mass on; the FIRST frame of each demo (phase == 0.0) exists only
    to satisfy the phase-convention invariant and always gets EXACTLY ZERO probability in these tests,
    so it never perturbs a weighted mean (0 * anything == 0) -- lets a test solve for the exact
    resulting anchor by hand.

    Returns (nodes, ctx, idx0, idx1) -- the node indices carrying q0/q1."""
    demos = [_fixed_demo(q0), _fixed_demo(q1)]
    nodes = _declare_ee_base(NodeTable.from_demos(demos, task_names=["taskA"], coarsen=1))
    ctx = ReadoutContext(nodes=nodes, topm=topm)
    return nodes, ctx, 1, 3


# the two arms this repo deploys (module docstring) -- checked directly against ARMS (matches
# schema.READOUT_ARMS).
LIVE_ARMS = ("euc_raw", "basin")


# ======================================================================================================
# registry
# ======================================================================================================
def test_registry_contains_all_arms():
    assert set(ARMS) == set(LIVE_ARMS)
    assert len(ARMS) == len(LIVE_ARMS)


def test_make_readout_unknown_name_raises_helpful_error():
    with pytest.raises(KeyError) as exc:
        make_readout("not_a_real_arm")
    msg = str(exc.value)
    for name in LIVE_ARMS:
        assert name in msg


def test_arms_are_stateless_singletons():
    assert make_readout("euc_raw") is make_readout("euc_raw") is ARMS["euc_raw"]
    assert make_readout("basin") is make_readout("basin") is ARMS["basin"]


# ======================================================================================================
# well-formed results
# ======================================================================================================
def test_well_formed_result_euc_raw():
    nodes, ctx = _flat_context()
    p = np.full(len(nodes), 1.0 / len(nodes))
    res = make_readout("euc_raw")(p, ctx)
    assert res.node is None
    assert res.q_star.shape == (D,)
    assert res.q_seg.shape == (ctx.seg_k, D)
    assert res.qdot_seg.shape == (ctx.seg_k, D)
    assert res.ee_seg.shape == (ctx.seg_k, 6)
    assert np.isfinite(res.q_star).all()
    assert np.isfinite(res.q_seg).all() and np.isfinite(res.qdot_seg).all()
    assert np.isfinite(res.ee_seg).all()
    assert np.isfinite(res.off_manifold)


def test_well_formed_result_basin():
    nodes, ctx = _flat_context()
    p = np.full(len(nodes), 1.0 / len(nodes))
    ctx.q_now = np.asarray(nodes.q[0], dtype=np.float64)
    ctx.basin = BasinGeometry(r=np.full((nodes.n_tasks, schema.NBINS_ALIGN), 0.5, dtype=np.float64), h=np.full((nodes.n_tasks, schema.NBINS_ALIGN), 0.1, dtype=np.float64))
    ctx.task_id = 0

    res = make_readout("basin")(p, ctx)
    assert res.node is None
    assert res.q_star.shape == (D,)
    assert res.q_seg.shape == (ctx.seg_k, D)
    assert res.qdot_seg.shape == (ctx.seg_k, D)
    assert res.ee_seg.shape == (ctx.seg_k, 6)
    assert np.isfinite(res.q_star).all()
    assert np.isfinite(res.q_seg).all() and np.isfinite(res.qdot_seg).all()
    assert np.isfinite(res.ee_seg).all()
    assert np.isfinite(res.off_manifold)


# ======================================================================================================
# euc_raw does not snap
# ======================================================================================================
def test_euc_raw_does_not_snap_on_spread_mass():
    nodes, ctx, tip0, tip1 = _fork_context()
    n = len(nodes)
    p = np.full(n, 1e-9)
    p[tip0] = 0.5
    p[tip1] = 0.5
    p = p / p.sum()

    res_raw = make_readout("euc_raw")(p, ctx)
    assert res_raw.node is None
    q_all = np.asarray(nodes.q, dtype=np.float64)
    dists = np.linalg.norm(q_all - res_raw.q_star[None, :], axis=1)
    assert dists.min() > 1e-6, "euc_raw's q_star must not coincide with any single node's q"
    assert res_raw.off_manifold > 1e-6


def test_euc_raw_concentrated_mass_recovers_the_node():
    """When p concentrates entirely on one node, the p-weighted mean IS that node's own q (up
    to the tiny floor mass on every other node) -- the free-form blend degenerates to the single-strand
    answer, same as any other aggregation would under concentrated mass."""
    nodes, ctx, tip0, _ = _fork_context()
    p = _concentrated_p(tip0, len(nodes))

    res = make_readout("euc_raw")(p, ctx)
    assert res.node is None
    np.testing.assert_allclose(res.q_star, np.asarray(nodes.q[tip0], dtype=np.float64), atol=1e-6)
    assert res.depth == pytest.approx(float(nodes.phase[tip0]), abs=1e-6)


# ======================================================================================================
# ee_seg — the same posterior, the same support, the same weights as q_seg
# ======================================================================================================
def test_ee_seg_carries_the_same_posterior_weights_as_q_seg():
    """The ee track is blended by the identical p-weighted tensordot q_seg gets, over the identical
    truncated support: with ee built as q's first 3 columns twice, the two must agree exactly (not
    approximately) on mass spread across two demos."""
    nodes, ctx, tip0, tip1 = _fork_context()
    p = np.full(len(nodes), 1e-9)
    p[tip0], p[tip1] = 0.5, 0.5
    p = p / p.sum()

    res = make_readout("euc_raw")(p, ctx)
    np.testing.assert_array_equal(res.ee_seg[:, :3], res.q_seg[:, :3])
    np.testing.assert_array_equal(res.ee_seg[:, 3:], res.q_seg[:, :3])


# ======================================================================================================
# weighted_topk truncation honoured
# ======================================================================================================
def test_weighted_topk_truncation_honoured_euc_raw():
    nodes, ctx, tip0, tip1 = _fork_context()
    n = len(nodes)
    p = np.full(n, 1e-9)
    p[tip0] = 0.5
    p[tip1] = 0.5
    p = p / p.sum()

    res = make_readout("euc_raw")(p, ctx)
    assert res.conf >= 0.99
    assert TOPK_K_MIN <= res.info["K"] <= min(ctx.topm, TOPK_K_MAX)


# ======================================================================================================
# determinism
# ======================================================================================================
def test_determinism_both_arms():
    nodes, ctx, tip0, tip1 = _fork_context()
    ctx.q_now = np.asarray(nodes.q[0], dtype=np.float64)
    ctx.basin = BasinGeometry(r=np.full((nodes.n_tasks, schema.NBINS_ALIGN), 0.5, dtype=np.float64), h=np.full((nodes.n_tasks, schema.NBINS_ALIGN), 0.1, dtype=np.float64))
    ctx.task_id = 0
    n = len(nodes)
    p = np.full(n, 1e-9)
    p[tip0] = 0.5
    p[tip1] = 0.5
    p = p / p.sum()

    for name in LIVE_ARMS:
        arm = make_readout(name)
        r1 = arm(p, ctx)
        r2 = arm(p, ctx)
        assert r1.node == r2.node
        assert r1.depth == r2.depth
        assert r1.conf == r2.conf
        assert r1.off_manifold == r2.off_manifold
        np.testing.assert_array_equal(r1.q_star, r2.q_star)
        np.testing.assert_array_equal(r1.q_seg, r2.q_seg)
        np.testing.assert_array_equal(r1.qdot_seg, r2.qdot_seg)


# ======================================================================================================
# basin -- minimal projection into the corpus-derived good-enough region (see BasinReadout's docstring)
# ======================================================================================================
def test_basin_registered_and_resolves():
    assert "basin" in LIVE_ARMS
    assert make_readout("basin") is ARMS["basin"]


def test_basin_requires_q_now():
    nodes, ctx = _flat_context()
    p = _concentrated_p(0, len(nodes))
    ctx.basin = BasinGeometry(r=np.full((nodes.n_tasks, schema.NBINS_ALIGN), 0.5, dtype=np.float64), h=np.full((nodes.n_tasks, schema.NBINS_ALIGN), 0.1, dtype=np.float64))
    ctx.task_id = 0
    with pytest.raises(ValueError, match="ctx.q_now"):
        make_readout("basin")(p, ctx)


def test_basin_requires_corpus_geometry():
    """Missing basin_r/basin_h/task_id must raise a message pointing at the missing corpus geometry --
    this arm has no measurement of its own to fall back to (BasinReadout's docstring)."""
    nodes, ctx = _flat_context()
    p = _concentrated_p(0, len(nodes))
    ctx.q_now = np.zeros(D, dtype=np.float64)
    with pytest.raises(ValueError, match="ctx.basin"):
        make_readout("basin")(p, ctx)


def test_basin_exact_no_op_when_inside():
    """||Delta|| <= r_eff: q_star must equal q_now EXACTLY (assert_array_equal, not approx) and
    info['no_op'] must be True -- the arm takes the q_now.copy() branch directly rather than relying on
    the projection formula to happen to land on q_now (BasinReadout's docstring, "THE MECHANISM")."""
    q0 = np.zeros(D, dtype=np.float64)
    q1 = np.full(D, 1.0, dtype=np.float64)
    nodes, ctx, idx0, idx1 = _two_node_context(q0, q1)
    p = np.zeros(len(nodes), dtype=np.float64)
    p[idx0], p[idx1] = 0.5, 0.5

    ctx.q_now = q0.copy()      # mu is ~0.5*ones(D); a huge r_eff certifies "inside" regardless
    ctx.basin = BasinGeometry(r=np.full((nodes.n_tasks, schema.NBINS_ALIGN), 10.0, dtype=np.float64), h=np.full((nodes.n_tasks, schema.NBINS_ALIGN), 1.0, dtype=np.float64))
    ctx.task_id = 0

    res = make_readout("basin")(p, ctx)
    assert res.node is None
    assert res.info["no_op"] is True
    assert res.info["moved_norm"] == 0.0
    np.testing.assert_array_equal(res.q_star, ctx.q_now)


def test_basin_outside_lands_on_segment_at_r_eff_from_anchor():
    """||Delta|| > r_eff: q_star must be the point on segment q_now -> q_b sitting EXACTLY r_eff short
    of q_b -- a closed-form check, not merely "moved some amount" (BasinReadout's docstring, "THE
    MECHANISM"). q_b is read off euc_raw's q_star on the SAME (p, ctx.topm), since both arms compute the
    identical p-weighted anchor."""
    q0 = np.zeros(D, dtype=np.float64)
    q1 = np.full(D, 10.0, dtype=np.float64)
    nodes, ctx, idx0, idx1 = _two_node_context(q0, q1)
    p = np.zeros(len(nodes), dtype=np.float64)
    p[idx0], p[idx1] = 0.5, 0.5

    q_b = make_readout("euc_raw")(p, ctx).q_star
    ctx.q_now = q0.copy()
    ctx.basin = BasinGeometry(r=np.full((nodes.n_tasks, schema.NBINS_ALIGN), 1.0, dtype=np.float64), h=np.full((nodes.n_tasks, schema.NBINS_ALIGN), 0.2, dtype=np.float64))
    ctx.task_id = 0

    res = make_readout("basin")(p, ctx)
    r_eff = 1.0 - 0.2
    assert res.node is None
    assert res.info["no_op"] is False
    assert res.info["r_eff"] == pytest.approx(r_eff)

    dist_to_anchor = float(np.linalg.norm(res.q_star - q_b))
    assert dist_to_anchor == pytest.approx(r_eff, abs=1e-9)

    direction = q_b - ctx.q_now
    diff = res.q_star - ctx.q_now
    cos = float(np.dot(diff, direction) / (np.linalg.norm(diff) * np.linalg.norm(direction)))
    assert cos == pytest.approx(1.0, abs=1e-9), "q_star must stay ON the segment q_now -> q_b"


def test_basin_segment_is_translated_not_shrunk():
    """q_seg* = q_seg_blend + (q_star - q_b), qdot unchanged -- the translation this arm performs
    (BasinReadout's docstring, "THE SEGMENT IS TRANSLATED, NOT SHRUNK"). ee_seg is deliberately
    exempt: the blend's alpha supplies the magnitude, so a basin projection on top double-corrects."""
    q0 = np.zeros(D, dtype=np.float64)
    q1 = np.full(D, 10.0, dtype=np.float64)
    nodes, ctx, idx0, idx1 = _two_node_context(q0, q1)
    p = np.zeros(len(nodes), dtype=np.float64)
    p[idx0], p[idx1] = 0.5, 0.5

    res_raw = make_readout("euc_raw")(p, ctx)
    ctx.q_now = q0.copy()
    ctx.basin = BasinGeometry(r=np.full((nodes.n_tasks, schema.NBINS_ALIGN), 1.0, dtype=np.float64), h=np.full((nodes.n_tasks, schema.NBINS_ALIGN), 0.2, dtype=np.float64))
    ctx.task_id = 0

    res = make_readout("basin")(p, ctx)
    expected_q_seg = res_raw.q_seg + (res.q_star - res_raw.q_star)[None, :]
    np.testing.assert_allclose(res.q_seg, expected_q_seg, atol=1e-9)
    np.testing.assert_allclose(res.qdot_seg, res_raw.qdot_seg, atol=1e-9)
    assert not np.allclose(res.q_star, res_raw.q_star)      # sanity: a translation really happened
    np.testing.assert_array_equal(res.ee_seg, res_raw.ee_seg)


# ======================================================================================================
# the differentiable torch expectation — must optimise the SAME quantity the numpy arms deploy
# ======================================================================================================
def test_ee_segment_table_rows_match_node_table_including_demo_boundaries():
    """Every row of the materialised table must be NodeTable.ee_segment exactly -- same clip at
    raw_ptr[owner+1], same edge padding. The final nodes of each demo are the ones that clip, so the
    test asserts at least one row really was padded rather than passing vacuously."""
    nodes, ctx = _flat_context()
    table = ee_segment_table(nodes, ctx.seg_k)
    assert table.shape == (len(nodes), ctx.seg_k, 6)

    for v in range(len(nodes)):
        np.testing.assert_array_equal(
            table[v].double().numpy(), nodes.ee_segment(v, ctx.seg_k)
        )
    padded = [v for v in range(len(nodes))
              if (nodes.raw_ptr[nodes.owner[v] + 1] - nodes.t_raw[v]) < ctx.seg_k]
    assert padded, "no node's segment was clipped: the boundary case is untested"


def test_expected_ee_segment_matches_the_numpy_aggregate():
    """The torch expectation and the deployed numpy _aggregate must be the SAME quantity, or Stage 2
    optimises something the rollout never computes. p is built so weighted_topk's 0.99-mass support
    is EXACTLY K nodes (asserted via info['K']), which is the only condition under which the two
    differ -- the torch arm uses a fixed k by design (expected_ee_segment's Note)."""
    k = 10
    nodes, ctx = _flat_context()
    # Distinct masses, so no sort/topk tie-break can pick a different support; the top k sum to
    # 0.9995 while the top k-1 sum to 0.9000, bracketing weighted_topk's 0.99 target at exactly k.
    p = np.full(len(nodes), 1e-9)
    p[:k] = 0.0995 + 1e-4 * np.arange(k)
    p = p / p.sum()

    res = make_readout("euc_raw")(p, ctx)
    assert res.info["K"] == k, "supports differ; the comparison below would be meaningless"

    got = expected_ee_segment(torch.as_tensor(p), ee_segment_table(nodes, ctx.seg_k), k=k)
    np.testing.assert_allclose(got.numpy(), res.ee_seg, rtol=1e-9, atol=1e-12)


def test_expected_ee_segment_gradient_reaches_p():
    """The whole point of the torch arm: an action-space loss must be able to move p. The topk
    INDICES are a hard selection, so nodes outside the support correctly get exactly zero grad."""
    v, seg_k, k = 6, 3, 3
    table = torch.arange(v * seg_k * 6, dtype=torch.float64).reshape(v, seg_k, 6)
    p = torch.tensor([0.30, 0.25, 0.20, 0.13, 0.08, 0.04], dtype=torch.float64, requires_grad=True)

    assert torch.autograd.gradcheck(lambda x: expected_ee_segment(x, table, k=k), (p,))

    expected_ee_segment(p, table, k=k).sum().backward()
    assert p.grad[:k].abs().sum() > 0
    assert torch.count_nonzero(p.grad[k:]) == 0