onf-c2/src/onf/graph/core/geometry.py

"""Geometry and shared train/deploy plumbing for the demonstration graph.

TRAIN/DEPLOY PARITY CONTRACT: velocity (finite_diff_vel), history padding
(edge_pad_hist) and off-manifold distance (off_manifold_distance) each have EXACTLY
ONE implementation, here. Independent reimplementations previously diverged (zero-pad vs
edge-extrapolate on row 0), producing different abstain logits for an identical window. Callers
route through these; never add a second path.
"""

from __future__ import annotations

import hashlib
import os
import warnings
from dataclasses import dataclass
from typing import Any, ClassVar, Optional, Protocol, runtime_checkable

import numpy as np
import torch
from scipy.special import expit

from onf.device import resolve_device
from onf.graph.core import schema
from onf.graph.core.edges import EdgeSet
from onf.graph.core.nodes import NodeTable

__all__ = [
    "NeighborIndex",
    "PaddedHistory",
    "GeometricField",
    "CleanlinessField",
    "CleanlinessConfig",
    "CleanlinessScorer",
    "finite_diff_vel",
    "edge_pad_hist",
    "off_manifold_distance",
    "off_manifold_batch",
    "leave_one_demo_out_floors",
    "graph_hash",
    "load_cleanliness_field",
    "CONSTANTS_JSON",
]

# C0's output, written beside g_nodes.npz. Hashed into graph_hash: a corpus recalibrated under a
# different recipe is a different graph even when the node table is byte-identical.
CONSTANTS_JSON = "constants.json"

# Fields hashed into graph_hash. w is in the list: under C2 it is a LEARNED column, not a pure
# function of the geometry, so a geometry-only hash would call two different graphs the same one.
_HASH_NODE_FIELDS = ("q", "qdot", "owner", "task_id", "stage", "t_idx", "phase")
_HASH_EDGE_FIELDS = ("src", "dst", "rel", "indptr", "w")
_HASH_VERSION = b"graph_hash_v2"
_HASH_ABSENT = b"<absent>"

# Floor under every user-supplied divisor, so a degenerate 0 cannot turn a weight into inf/nan.
_EPS = 1e-8


class NeighborIndex:
    """Blocked nearest-neighbour queries against a fixed corpus of joint configs.

    Uploads the corpus once and reuses it, so a caller that asks many times pays the transfer once.
    Optional per-row owner ids turn a plain NN query into the leave-one-demo-out floor: corpus rows
    sharing the query row's own demo are masked to inf.

    device is resolved once here and deliberately NOT plumbed through the module-level helpers:
    CPU and CUDA disagree by up to ~1e-4 rad elementwise, enough to move a training label.
    """

    def __init__(
        self,
        corpus: np.ndarray,
        corpus_owner: Optional[np.ndarray] = None,
        block_size: Optional[int] = None,
    ) -> None:
        """
        Args:
            corpus: [M, D] reference configs; cast to float32.
            corpus_owner: Optional [M] demo ids, required for owner-excluded queries.
            block_size: Query rows per cdist call; defaults to schema.CDIST_BS.
        """
        self._device = resolve_device(None)
        self._block_size = int(block_size if block_size is not None else schema.CDIST_BS)
        self._corpus = torch.as_tensor(
            np.asarray(corpus, dtype=np.float32), device=self._device
        )
        self._owner: Optional[torch.Tensor] = None
        if corpus_owner is not None:
            self._owner = torch.as_tensor(
                np.asarray(corpus_owner, dtype=np.int64), device=self._device
            )

    @classmethod
    def from_nodes(cls, nodes: NodeTable, with_owner: bool = False) -> "NeighborIndex":
        """Build an index over a graph's own node configs.

        Args:
            nodes: Source node table.
            with_owner: Attach nodes.owner to enable owner-excluded queries.

        Returns:
            A ready-to-query index.
        """
        owner = np.asarray(nodes.owner, dtype=np.int64) if with_owner else None
        return cls(np.asarray(nodes.q, dtype=np.float32), owner)

    def __len__(self) -> int:
        return int(self._corpus.shape[0])

    def min_distance(
        self, points: np.ndarray, point_owner: Optional[np.ndarray] = None
    ) -> np.ndarray:
        """L2 distance from each query row to its nearest corpus row.

        Args:
            points: [N, D] (or [D]) query configs; cast to float32.
            point_owner: Optional [N] demo ids; corpus rows with a matching owner are excluded
                from that row's minimum.

        Returns:
            [N] float64 distances, in radians.

        Raises:
            ValueError: point_owner given but the index carries no corpus owners.
        """
        query = np.atleast_2d(np.asarray(points, dtype=np.float32))
        n = query.shape[0]
        if n == 0:
            return np.zeros(0, dtype=np.float64)
        if point_owner is not None and self._owner is None:
            raise ValueError("owner-excluded query requires a corpus_owner-bearing index")

        query_t = torch.as_tensor(query, device=self._device)
        owner_t: Optional[torch.Tensor] = None
        if point_owner is not None:
            owner_t = torch.as_tensor(
                np.asarray(point_owner, dtype=np.int64), device=self._device
            )

        out = np.empty(n, dtype=np.float64)
        for start in range(0, n, self._block_size):
            stop = start + self._block_size
            dist = torch.cdist(query_t[start:stop], self._corpus)
            if owner_t is not None:
                excluded = self._owner[None, :] == owner_t[start:stop, None]
                dist = dist.masked_fill(excluded, float("inf"))
            out[start:stop] = dist.min(dim=1).values.double().cpu().numpy()
        return out


def finite_diff_vel(q_hist: np.ndarray) -> np.ndarray:
    """vh[t] = q_hist[t] - q_hist[t-1], edge-EXTRAPOLATED at t=0 with vh[1].

    Never zero-padded, and derived from q_hist alone (never a separate qdot channel), so a
    window built at train time and the same raw content seen live produce byte-identical velocities
    — and therefore an identical e_Q and abstain decision. Single source of truth: the two
    former implementations disagreed on row 0 and silently split the abstain logits.

    Takes any leading batch shape ([T, D] or [N, T, D]); the difference is along axis -2.
    """
    q_hist = np.asarray(q_hist, dtype=np.float32)
    v = np.zeros_like(q_hist)
    t = q_hist.shape[-2]
    if t > 1:
        v[..., 1:, :] = q_hist[..., 1:, :] - q_hist[..., :-1, :]
        v[..., 0, :] = v[..., 1, :]
    return v


def off_manifold_distance(nodes: NodeTable, q_last: np.ndarray) -> float:
    """Nearest-node L2 distance from one joint config to any node, with no owner exclusion.

    DELIBERATELY NOT routed through NeighborIndex, unlike its batched sibling: this is a
    float64 numpy reduction while the index replays torch.cdist in float32 on whatever device
    is available. The two agree to ~1e-6 rad — fine for a training label, not fine here, since this
    value is pinned bit-exactly in tests/golden/parity.json (where_off_manifold_hex).

    Args:
        nodes: Graph node table.
        q_last: [D] joint config.

    Returns:
        The distance in radians.
    """
    q_last = np.asarray(q_last, dtype=np.float64).reshape(-1)
    return float(np.min(np.linalg.norm(nodes.q.astype(np.float64) - q_last[None, :], axis=1)))


def off_manifold_batch(nodes: NodeTable, q_last: np.ndarray) -> np.ndarray:
    """Vectorised off_manifold_distance over q_last [N, D] — same definition, one
    torch.cdist pass instead of N numpy calls.

    Args:
        nodes: Graph node table.
        q_last: [N, D] joint configs.

    Returns:
        [N] float64 distances, in radians.
    """
    return NeighborIndex.from_nodes(nodes).min_distance(q_last)


@dataclass(frozen=True, slots=True)
class PaddedHistory:
    """edge_pad_hist's result. Unpacks and subscripts as the (padded, was_padded) tuple
    it replaced, so callers writing arr, flag = edge_pad_hist(...) are unchanged.

    Attributes:
        window: [hist_len, ...] right-aligned array.
        was_padded: True when the input was shorter than hist_len and was edge-padded.
    """

    window: np.ndarray
    was_padded: bool

    def __iter__(self):
        """Yields window then was_padded."""
        yield self.window
        yield self.was_padded

    def __getitem__(self, i: int):
        """[0] is window, [1] is was_padded."""
        return (self.window, self.was_padded)[i]


def edge_pad_hist(arr: np.ndarray, hist_len: int) -> PaddedHistory:
    """Right-align arr to exactly hist_len rows: truncate to the most recent rows if longer,
    edge-pad (repeat the oldest row at the front) if shorter.

    Single source of truth for both the deploy and the train path, so a variable-length training
    window and a variable-length live q_hist are padded identically.

    Args:
        arr: [T, ...] input, oldest row first.
        hist_len: Target row count.

    Returns:
        The normalised window and whether padding occurred.
    """
    arr = np.asarray(arr)
    n = arr.shape[0]
    padded = n < hist_len
    if padded:
        pad_n = hist_len - n
        arr = np.concatenate([np.repeat(arr[:1], pad_n, axis=0), arr], axis=0)
    elif n > hist_len:
        arr = arr[-hist_len:]
    return PaddedHistory(arr, padded)


def graph_hash(
    nodes: NodeTable, edges: EdgeSet, *, constants: Any = None, c2_state: Any = None,
) -> str:
    """Deterministic content fingerprint of one (graph, constants, C2) triple — v2.

        graph_hash_v2 = SHA-256(node+edge geometry || constants.json || C2 checkpoint)

    The exact graph a retrieval head was trained on, so a head cannot be loaded against a graph it
    was not fit to. Both new terms are load-bearing: C0's constants and C2's weights change what the
    graph MEANS while leaving its geometry byte-identical, and the resulting failure is a foreign
    head returning confident, correctly-shaped, wrong answers. An absent term hashes as a marker,
    never as nothing, so a v1 stamp can never accidentally equal a v2 one.

    Args:
        nodes: The node table.
        edges: Its edge set.
        constants: C0's constants.json, as a path or as raw bytes/str; None when the graph shipped
            none.
        c2_state: The C2 checkpoint, as a mapping of key to array (see
            onf.graph.net.gnn.GraphRetrieverNet.c2_state); None when the head carries no C2.

    Returns:
        The hex digest.
    """
    h = hashlib.sha256()
    h.update(_HASH_VERSION)
    h.update(f"V={len(nodes)}|E={len(edges.src)}|dim={nodes.dim}".encode())
    for name in _HASH_NODE_FIELDS:
        h.update(np.ascontiguousarray(getattr(nodes, name)).tobytes())
    for name in _HASH_EDGE_FIELDS:
        h.update(np.ascontiguousarray(getattr(edges, name)).tobytes())
    h.update(b"|constants|")
    h.update(_constants_bytes(constants))
    h.update(b"|c2|")
    h.update(_c2_bytes(c2_state))
    return h.hexdigest()


def _constants_bytes(constants: Any) -> bytes:
    """C0's constants.json as bytes: a path is read, bytes/str pass through, None is the marker."""
    if constants is None:
        return _HASH_ABSENT
    if isinstance(constants, (bytes, bytearray)):
        return bytes(constants)
    if isinstance(constants, str) and not os.path.exists(constants):
        return constants.encode()
    path = os.fspath(constants)
    return _HASH_ABSENT if not os.path.exists(path) else open(path, "rb").read()


def _c2_bytes(c2_state: Any) -> bytes:
    """The C2 checkpoint as bytes: keys in sorted order, each followed by its array's raw buffer."""
    if c2_state is None:
        return _HASH_ABSENT
    parts: list[bytes] = []
    for key in sorted(c2_state):
        parts.append(key.encode())
        parts.append(np.ascontiguousarray(c2_state[key]).tobytes())
    return _HASH_ABSENT if not parts else b"|".join(parts)


def leave_one_demo_out_floors(
    points: np.ndarray, point_owner: np.ndarray, corpus: np.ndarray, corpus_owner: np.ndarray,
) -> np.ndarray:
    """Nearest-neighbour distance from each point to the corpus, excluding rows from that point's
    OWN demo. onf.graph.core.corpus keeps a numpy replay of this on purpose.

    Args:
        points: [N, D] query configs.
        point_owner: [N] demo ids for points.
        corpus: [M, D] reference configs.
        corpus_owner: [M] demo ids for corpus.

    Returns:
        [N] float64 floors, in radians.
    """
    return NeighborIndex(corpus, corpus_owner).min_distance(points, point_owner)


# ---- cleanliness: the WEIGHT half of QueryEncoder's pooling -------------------------------------
# The rule must MATCH between train and deploy -- a head trained under one pooling and evaluated
# under the other is a silent train/deploy mismatch, not an ablation arm. So ONE CleanlinessConfig
# is built at the composition root and injected into both paths; nothing downstream re-reads env.


@runtime_checkable
class CleanlinessField(Protocol):
    """The duck type CleanlinessScorer reads. Pins GeometricField and the trained
    onf.field.field.ONFField (not importable here at module scope) to one surface."""

    def f_value(self, q0: np.ndarray) -> float:
        """Field distance f(q) at a single [D] config."""
        ...


@dataclass(frozen=True, slots=True)
class CleanlinessConfig:
    """Which cleanliness field a run uses, and how sharply it weights.

    Attributes:
        use_geometric: Use the weight-free GeometricField instead of a trained field.
        geo_scale: Distance scale in radians for the geometric field.
        temp: Sigmoid temperature applied to field distances.
    """

    use_geometric: bool = False
    geo_scale: float = schema.GEO_SCALE
    temp: float = 1.0

    @classmethod
    def from_env(cls, temp: float = 1.0) -> "CleanlinessConfig":
        """Read ONF_CLEANLINESS/ONF_GEO_SCALE — the ONE place in the codebase that does.

        Args:
            temp: Sigmoid temperature; not environment-driven.

        Returns:
            The config this process's environment selects.
        """
        return cls(
            use_geometric=os.environ.get("ONF_CLEANLINESS", "field").lower() == "geo",
            geo_scale=float(os.environ.get("ONF_GEO_SCALE", schema.GEO_SCALE)),
            temp=temp,
        )


class GeometricField:
    """Weight-free stand-in for onf.field.field.ONFField:
    f(q) = ||q - nearest node|| / scale.

    scale sets how fast w = sigmoid(-f) falls off; calibrated picks it so an arm
    using this field is comparable to one using a trained field rather than differing by a gain.

    Args:
        q_nodes: [M, D] node configs.
        scale: Distance normaliser in radians.
    """

    def __init__(self, q_nodes: np.ndarray, scale: float):
        self.q = np.asarray(q_nodes, dtype=np.float64)
        self.scale = float(scale)

    @classmethod
    def calibrated(cls, nodes: NodeTable, median_f: float, sample: int = 4096) -> "GeometricField":
        """Build with scale set so the median f over sampled leave-one-demo-out node
        distances equals median_f (a trained field's own median on the same graph).

        Args:
            nodes: Graph node table.
            median_f: Target median field distance.
            sample: Number of nodes to sample.

        Returns:
            A calibrated field.
        """
        q = np.asarray(nodes.q, dtype=np.float64)
        owner = np.asarray(nodes.owner, dtype=np.int64)
        idx = np.linspace(0, len(q) - 1, min(sample, len(q))).astype(np.int64)
        d = leave_one_demo_out_floors(q[idx], owner[idx], q, owner)
        return cls(q, float(np.median(d)) / max(float(median_f), _EPS))

    def f_value(self, q0: np.ndarray) -> float:
        """Field distance f(q) at a single [D] config, in units of scale."""
        q0 = np.asarray(q0, dtype=np.float64).reshape(-1)
        return float(np.min(np.linalg.norm(self.q - q0[None, :], axis=1))) / self.scale


def load_cleanliness_field(
    nodes: NodeTable, config: CleanlinessConfig, field_dir: Any = None
) -> CleanlinessField | None:
    """Resolve the field object a CleanlinessScorer should read.

    Args:
        nodes: Graph node table, used when the geometric field is selected.
        config: Cleanliness settings; see CleanlinessConfig.from_env.
        field_dir: Directory holding a trained field; ignored under config.use_geometric.

    Returns:
        A GeometricField, a trained ONFField, or None when neither is configured or
        the directory holds no field (uniform weights — not an error).
    """
    if config.use_geometric:
        return GeometricField(nodes.q, config.geo_scale)
    if field_dir is None:
        return None
    try:
        from onf.field.field import ONFField

        return ONFField.load(field_dir)
    except (FileNotFoundError, OSError):
        return None


class CleanlinessScorer:
    """Computes the per-step weights w_i = sigmoid(-f(q_i) / temp) that
    onf.graph.net.modules.QueryEncoder pools a query window by.

    Holds the field and temperature instead of taking them per call, so a control loop builds one
    scorer at startup and the rule cannot drift between steps. Falls back to uniform weights when no
    field is available, warning once per process — the fallback is a property of the run.

    Args:
        field: A trained or geometric field, or None for uniform weights.
        temp: Sigmoid temperature.
    """

    _uniform_warning_emitted: ClassVar[bool] = False

    def __init__(self, field: CleanlinessField | None, temp: float = 1.0) -> None:
        self.field = field
        self.temp = max(float(temp), _EPS)
        if field is None:
            self._warn_uniform_once()

    @classmethod
    def from_graph(
        cls, nodes: NodeTable, config: CleanlinessConfig, field_dir: Any = None
    ) -> "CleanlinessScorer":
        """Build the scorer a configuration implies.

        Args:
            nodes: Graph node table, used when the geometric field is selected.
            config: Cleanliness settings.
            field_dir: Directory holding a trained field; ignored under config.use_geometric.

        Returns:
            A ready scorer.
        """
        return cls(load_cleanliness_field(nodes, config, field_dir), config.temp)

    @classmethod
    def reset_uniform_warning(cls) -> None:
        """Re-arm the once-per-process uniform-weight warning. For tests."""
        cls._uniform_warning_emitted = False

    @classmethod
    def _warn_uniform_once(cls) -> None:
        """Emit the uniform-weight warning unless it has already fired this process."""
        if cls._uniform_warning_emitted:
            return
        cls._uniform_warning_emitted = True
        warnings.warn(
            "no ONF field available for this suite -- falling back to uniform window weights (the "
            "query pools as a plain average instead of favoring the clean steps)",
            stacklevel=3,
        )

    @property
    def is_uniform(self) -> bool:
        """True when no field backs this scorer, so every weight is 1.0."""
        return self.field is None

    def weights(self, q_window: np.ndarray) -> np.ndarray:
        """Cleanliness weight per step of one window.

        Args:
            q_window: [T, D] joint configs.

        Returns:
            [T] float64 weights in (0, 1), or ones when is_uniform.
        """
        q_window = np.atleast_2d(np.asarray(q_window, dtype=np.float64))
        if self.field is None:
            return np.ones(q_window.shape[0], dtype=np.float64)
        f = np.array([self.field.f_value(q) for q in q_window], dtype=np.float64)
        return expit(-f / self.temp)

    def weights_batch(self, q_hist: np.ndarray) -> np.ndarray:
        """Cleanliness weights for a stacked query set.

        A per-window loop on purpose: field's only guaranteed method is the scalar f_value,
        so there is no batched surface. Runs once at dataset-build time, not per epoch.

        Args:
            q_hist: [N, T, D] stacked query windows.

        Returns:
            [N, T] float32 weights.
        """
        return np.stack([self.weights(window) for window in q_hist]).astype(np.float32)