runtime/src/resynthesis/trauma_system.py

"""Tensor-native trauma system — the definitive HARD-WON KNOWLEDGE ledger.

OPERATOR DOCTRINE (additive-only -- governs every helper below): the 12b
parent/base is a FROZEN HOLLOW VOCAB SUBSTRATE (a tiny lego that just
provides tokenization); the REAL model is the 1T+ ADDITIVE machinery (NoNE
pages, experts, RBO, causal/trauma/MitM/coverage-pressure) that we train.
THIS module operates ONLY on the ADDITIVE arms (per-arm trauma scalars +
the additive knowledge ledger); it never touches/improves/retains the
frozen base -- the base is barely in the model and is never a trauma or
retention target here.

CANONICAL DOCTRINE (trauma = the definitive hard knowledge)
-----------------------------------------------------------
**Trauma IS the hard knowledge** — the definitive, hard-won knowledge the
model MUST learn (positive AND negative). It is the **ledger of what is
definitively known/not-yet-known**, and it anchors BOTH curriculum priority
AND preservation:

  * **POSITIVE trauma = PRESERVE**. Hard-won truths / abilities / knowledge
    the model MASTERED through difficulty. This is **definitive knowledge that
    must be PRESERVED** — the anti-forgetting anchor. Once hard-won (high
    confidence, survived verification), it is stamped ``positive-definitive``
    and feeds preservation/replay weight so training must never lose it.
    Surfaced via ``positive_definitive_preservation_weights``.

  * **NEGATIVE trauma = AVOID**. Hard anti-patterns / definitive failures the
    model must AVOID. This is the contrastive / anti signal. Stamped
    ``negative-definitive`` when a failure survives verification at high
    confidence, it feeds the contrastive/anti curriculum.

  * **DEFINITIVE = ANCHOR**. Both polarities are "definitive" once they are
    HIGH-confidence hard-won knowledge that SURVIVED VERIFICATION. Definitive
    arms get anchor status: top curriculum priority + preservation weight.
    ``mark_definitive`` is the ONLY stamp path, and it requires an explicit
    verification signal (a probe pass after struggle, a verified failure) — it
    is NEVER auto-stamped on noise. See ``DEFAULT_DEFINITIVE_CONFIDENCE``.

  * **MUST-LEARN CURRICULUM PRIORITY**. The definitive hard knowledge gets TOP
    learning priority. ``hard_knowledge_must_learn_priority`` ranks both
    polarities by ``trauma_level x definitive_confidence x (1 - coverage)`` —
    this is the signal coverage-pressure targeting + the MitM scaffold consume
    to point onto the RIGHT hard targets (positive-definitive to reinforce/
    preserve; negative-definitive to contrast/avoid).

LEGACY FRAMING (unchanged behavior, restated)
---------------------------------------------
Trauma is ALSO the **surface of what is hard** — the curriculum the hill-climb,
MITM bridge, and coverage pressure target. Anti-Thompson repels collapsed wrong
paths; trauma marks what remains definitively unlearned or definitively earned.

Adapted (not copied) from a prior external trauma stack into
resynthesis's tensor-native, schema-sealed idiom.

DESIGN INVARIANTS (resynthesis-native):
  - State is CHEAP: O(num_arms) scalars per lane; the additive knowledge ledger
    is a small O(num_arms) dict, off-by-default, never read by the hot path.
  - FAIL-OPEN everywhere — trauma never aborts a transaction.
  - Off-by-default via ``NNF_TRAUMA_SYSTEM=1``. The additive
    definitive-knowledge layer (binding / stamps / ledger) is ADDITIVE and
    default-empty: existing behavior is byte-identical when the new fields are
    unset. Off-by-default via ``NNF_TRAUMA_DEFINITIVE_LEDGER=1``.
  - Trauma nudges routing/verification targeting; it never peeks answers.
  - Definitive stamps require an explicit verification signal (never noise).
"""

from __future__ import annotations

import hashlib
from collections.abc import Callable
from typing import cast

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor

TRAUMA_SYSTEM_SCHEMA = "nnf.resynthesis.trauma_system.v1"
HARD_KNOWLEDGE_RECEIPT_SCHEMA = "nnf.resynthesis.hard_knowledge_surface.v1"

# --- DEFAULT_* scalars (resynthesis idiom; see anti_systems_bridge.py) --------
DEFAULT_TRAUMA_FAIL_DECAY = 0.90
"""Per-step exponential decay on the negative (fail) EMA. <1 so trauma fades
when an arm stops failing; matches the anti-systems bridge default family."""

DEFAULT_TRAUMA_SUCCESS_DECAY = 0.94
"""Per-step exponential decay on the positive (success) EMA. Higher than the
fail decay so successes persist longer than failures (tempered reinforcement)."""

DEFAULT_TRAUMA_PRESSURE_SCALE = 0.25
"""Scale of the negative repulsion bias added to the quantile router. Matches
``DEFAULT_ANTI_BIAS_SCALE`` so trauma is the same strength family as the
existing anti-Thompson nudge (additive, not dominating)."""

DEFAULT_TRAUMA_POSITIVE_SCALE = 0.15
"""Scale of the positive reinforcement bias. Smaller than the negative scale
on purpose: reinforcing collapse is more dangerous than repelling a bad arm,
so the positive path is more conservative (anti-collapse first)."""

DEFAULT_TRAUMA_COOLDOWN_STEPS = 8
"""Number of recent failure-laden steps that must clear before positive
reinforcement is allowed to fire. Adapted from the source trainer's
``adaptation_cooldown``; prevents reinforcing an arm that only just stopped
failing (would re-traumatize immediately)."""

DEFAULT_TRAUMA_HIGH_FAIL_THRESHOLD = 0.6
"""``fail_ema`` above which an arm is considered "actively traumatized" and
contributes to the cooldown gate (positive reinforcement globally suppressed
while any arm is this hot). Adapted from the source trainer's
``high_trauma_threshold``."""

DEFAULT_TRAUMA_MAX_FRACTION = 0.5
"""Anti-collapse cap: no single arm may capture more than this fraction of the
total positive-reinforcement mass. If one arm would dominate, the surplus is
redistributed. This is the tempered-reinforcement guard the operator asked
for ("avoid collapse to one arm")."""

DEFAULT_TRAUMA_EPS = 1.0e-8
"""Numerical floor for divisions (matches the eps family used elsewhere)."""

DEFAULT_HARD_KNOWLEDGE_FAIL_THRESHOLD = 0.35
"""Arms at/above this ``fail_ema`` (or ``peak_fail_ema``) are **negative hard
knowledge** — definitive gaps still missing from weights/pages."""

DEFAULT_HARD_WON_STRUGGLE_THRESHOLD = 0.25
"""Success on an arm whose ``peak_fail_ema`` reached here counts as struggle —
the win is **positive hard knowledge** (hard-won, not easy)."""

DEFAULT_HARD_WON_EMA_THRESHOLD = 0.20
"""Minimum ``hard_won_ema`` for an arm to appear on the positive hard-knowledge
surface (definitive earned capability)."""

# --- Definitive-knowledge ledger scalars (additive, off-by-default) ---------
#
# These gate the DEFINITIVE layer (knowledge binding, confidence stamps,
# preservation weights, must-learn priority, ledger receipt). All ADDITIVE:
# the existing trauma behavior is byte-identical when these are unused (the
# ledger fields default-empty). The layer is OFF-BY-DEFAULT and must be
# explicitly armed by the operator; nothing here fires on the training hot
# path unless a caller (e.g. an explicit verification boundary) invokes it.

DEFAULT_DEFINITIVE_CONFIDENCE = 0.85
"""Minimum confidence for a knowledge arm to be stamped ``definitive``.

Only HIGH-confidence, verification-survived knowledge becomes "definitive"
(the anchor class). Below this threshold the arm stays a candidate (hard
knowledge, but not yet an anchor). This is the doctrine gate: definitive
stamps require an explicit verification signal, never noise."""

DEFAULT_DEFINITIVE_CONFIDENCE_FLOOR = 0.0
"""Lower clamp on a definitive confidence stamp (a caller must not be able to
push a negative/NaN confidence into the ledger)."""

DEFAULT_DEFINITIVE_CONFIDENCE_CEIL = 1.0
"""Upper clamp on a definitive confidence stamp."""

DEFAULT_POSITIVE_DEFINITIVE_PRESERVATION_SCALE = 1.0
"""Scale of the per-arm preservation weight for positive-definitive knowledge.
A scalar multiplier on the normalized preservation surface so the caller can
temper how strongly hard-won positive knowledge is protected."""

DEFAULT_MUST_LEARN_PRIORITY_TOP_K = 64
"""Default ``top_k`` for the must-learn curriculum priority surface."""

DEFAULT_MUST_LEARN_COVERAGE_EPS = 1.0e-6
"""Numerical floor for the (1 - coverage) term in the must-learn priority so an
exactly-covered arm does not produce a 0/0 (it simply gets ~zero priority)."""

ENV_TRAUMA_DEFINITIVE_LEDGER = "NNF_TRAUMA_DEFINITIVE_LEDGER"
"""Env gate for the additive definitive-knowledge layer. Default off; the
operator arms it to persist/consume knowledge bindings + definitive stamps.
All public ledger functions FAIL-OPEN to empty when unset, so the hot path
(training loop, router) is unaffected."""

DEFINITIVE_POLARITY_POSITIVE = "positive"
"""Stamp polarity: positive-definitive (hard-won truth to PRESERVE)."""

DEFINITIVE_POLARITY_NEGATIVE = "negative"
"""Stamp polarity: negative-definitive (hard anti-pattern to AVOID)."""

DEFINITIVE_POLARITY_NONE = "none"
"""Stamp polarity: arm is hard knowledge but NOT yet an anchor (no
high-confidence verification has survived)."""

# --- Tensor-encoding helpers for the definitive-knowledge ledger -------------
#
# The model-facing in-memory state is now TENSOR-NATIVE: polarity is stored as
# an int8 per-arm tensor (0=none, +1=positive, -1=negative), confidence as a
# float per-arm tensor, and the knowledge-id binding as a long tensor of stable
# hash tokens (0 reserved for unbound). The legacy ``dict[int,str]`` /
# ``tuple[str,...]`` / ``dict[str,int]`` views are kept as cheap DERIVED
# properties (reconstructed from the tensors) so every existing caller, the
# read-only ledger receipt, and snapshot/restore stay byte-identical.
DEFINITIVE_POLARITY_CODE: dict[str, int] = {
    DEFINITIVE_POLARITY_NONE: 0,
    DEFINITIVE_POLARITY_POSITIVE: 1,
    DEFINITIVE_POLARITY_NEGATIVE: -1,
}
"""Map polarity string -> int8 tensor code (model-facing representation)."""

DEFINITIVE_POLARITY_DECODE: dict[int, str] = {
    code: name for name, code in DEFINITIVE_POLARITY_CODE.items()
}
"""Reverse map int8 tensor code -> polarity string (for the derived dict view)."""

DEFINITIVE_UNBOUND_TOKEN: int = 0
"""Reserved knowledge-id hash token meaning "this arm has no binding".
Non-zero so a never-bound arm is exactly distinguishable from a bound one
under any masked-reduction (mirrors the trauma fail/success EMA mask idiom)."""


def _knowledge_id_to_token(knowledge_id: str) -> int:
    """Stable non-zero 63-bit hash token for a knowledge id (0 = unbound).

    Uses BLAKE2b (8-byte digest) so the token is deterministic across processes
    and Python hash-seed changes (unlike the built-in ``hash``). The result is
    masked to 63 bits and forced non-zero (collisions re-hash) so the reserved
    ``DEFINITIVE_UNBOUND_TOKEN`` (0) is never produced for a real id. This is
    the model-facing integer id per arm; the reverse bookkeeping
    (token -> id string) lives in ``_knowledge_token_to_id`` purely for the
    derived dict/tuple views and the read-only ledger receipt.
    """

    kid = str(knowledge_id)
    if not kid.strip():
        return DEFINITIVE_UNBOUND_TOKEN
    digest = hashlib.blake2b(kid.encode("utf-8"), digest_size=8).digest()
    token = int.from_bytes(digest, "big") & 0x7FFFFFFFFFFFFFFF
    # Guarantee non-zero (0 is reserved for "unbound"); re-hash on the
    # astronomically unlikely zero collision.
    seed = 0
    while token == DEFINITIVE_UNBOUND_TOKEN and seed < 8:
        digest = hashlib.blake2b(
            (kid + "\x00" * seed).encode("utf-8"), digest_size=8
        ).digest()
        token = int.from_bytes(digest, "big") & 0x7FFFFFFFFFFFFFFF
        seed += 1
    if token == DEFINITIVE_UNBOUND_TOKEN:  # pragma: no cover - unreachable
        token = 1
    return token

DEFINITIVE_KNOWLEDGE_LEDGER_SCHEMA = (
    "nnf.resynthesis.definitive_knowledge_ledger.v1"
)
"""Schema constant for the read-only definitive-knowledge ledger receipt."""


class TensorTraumaState(nn.Module):
    """Per-arm trauma bank — the definitive hard-knowledge ledger (tensor-native).

    Each arm (page / expert / capability index) carries:

      * ``fail_ema`` — negative hard knowledge (gap not yet learned)
      * ``success_ema`` — recent success mass (routing reinforcement input)
      * ``peak_fail_ema`` — worst gap this arm ever hit (struggle marker)
      * ``hard_won_ema`` — positive hard knowledge (success AFTER struggle)
      * ``last_success_step`` / ``global_step`` — cooldown + ordering

    ADDITIVE definitive-knowledge ledger (off-by-default, default-empty) —
    now TENSOR-NATIVE in memory (operator directive: "everything needs to be a
    tensor on your enhancements"):
      * ``knowledge_id_tokens_t`` — ``[num_arms]`` long tensor of stable hash
        tokens (one per bound knowledge id, ``DEFINITIVE_UNBOUND_TOKEN`` = 0
        for unbound). This is the model-facing per-arm knowledge identity.
      * ``definitive_polarity_t`` — ``[num_arms]`` int8 tensor (0 = none,
        +1 = positive, -1 = negative). Updated only by ``mark_definitive``.
      * ``definitive_confidence_t`` — ``[num_arms]`` float tensor in [0, 1];
        0.0 until stamped. Key membership tracks ``definitive_polarity_t``.
      * ``_knowledge_token_to_id`` — small ``dict[int, str]`` reverse-lookup
        map (token -> id string) kept ONLY for the derived dict/tuple views
        and the read-only ledger receipt. Bookkeeping, never model-facing.

    The legacy ``arm_knowledge_ids`` / ``knowledge_to_arm`` /
    ``definitive_polarity`` / ``definitive_confidence`` fields are kept as
    cheap DERIVED properties (reconstructed from the tensors) so EVERY
    existing caller, the read-only ledger receipt, and snapshot/restore stay
    byte-identical. The legacy dict/tuple forms are now views, not the source
    of truth — the tensors are.

    State is CHEAP: O(num_arms) scalars, device-agnostic, JSONL-persistable.
    """

    fail_ema: Tensor
    success_ema: Tensor
    peak_fail_ema: Tensor
    hard_won_ema: Tensor
    last_success_step: Tensor
    global_step: Tensor
    knowledge_pro_ema: Tensor
    knowledge_anti_ema: Tensor
    behavior_pro_ema: Tensor
    behavior_anti_ema: Tensor
    verified_exposure_ema: Tensor
    evidence_confidence_ema: Tensor
    definitive_polarity_t: Tensor
    definitive_confidence_t: Tensor
    knowledge_id_tokens_t: Tensor

    # --- construction (signature identical to the legacy @dataclass) ---------

    def __init__(
        self,
        num_arms: int,
        fail_decay: float = DEFAULT_TRAUMA_FAIL_DECAY,
        success_decay: float = DEFAULT_TRAUMA_SUCCESS_DECAY,
        *,
        # Legacy keyword args (kept for snapshot/restore + any caller that
        # constructs with explicit bindings). Accepted as keyword-only so the
        # positional signature (num_arms, fail_decay, success_decay) is
        # unchanged and the constructor never has to disambiguate.
        arm_knowledge_ids: tuple[str, ...] = (),
        knowledge_to_arm: dict[str, int] | None = None,
        definitive_polarity: dict[int, str] | None = None,
        definitive_confidence: dict[int, float] | None = None,
    ) -> None:
        super().__init__()
        if num_arms < 1:
            raise ValueError("trauma state requires at least one arm")
        if not (0.0 < fail_decay <= 1.0):
            raise ValueError("fail_decay must be in (0, 1]")
        if not (0.0 < success_decay <= 1.0):
            raise ValueError("success_decay must be in (0, 1]")

        # Core trauma tensors (unchanged).
        self.num_arms = int(num_arms)
        self.fail_decay = float(fail_decay)
        self.success_decay = float(success_decay)
        self.register_buffer(
            "fail_ema",
            torch.zeros(self.num_arms, dtype=torch.float32),
            persistent=True,
        )
        self.register_buffer(
            "success_ema",
            torch.zeros(self.num_arms, dtype=torch.float32),
            persistent=True,
        )
        self.register_buffer(
            "peak_fail_ema",
            torch.zeros(self.num_arms, dtype=torch.float32),
            persistent=True,
        )
        self.register_buffer(
            "hard_won_ema",
            torch.zeros(self.num_arms, dtype=torch.float32),
            persistent=True,
        )
        self.register_buffer(
            "last_success_step",
            torch.full((self.num_arms,), -1, dtype=torch.long),
            persistent=True,
        )
        self.register_buffer(
            "global_step",
            torch.zeros((), dtype=torch.long),
            persistent=True,
        )
        # Verified evidence remains separated by semantic axis.  The combined
        # fail/success EMAs above retain the public compatibility surface,
        # while these buffers prove whether knowledge or behavior evidence
        # caused the pressure.  All are persistent model state and therefore
        # move with ``module.to(device)`` and survive exact cold reload.
        for buffer_name in (
            "knowledge_pro_ema",
            "knowledge_anti_ema",
            "behavior_pro_ema",
            "behavior_anti_ema",
            "verified_exposure_ema",
            "evidence_confidence_ema",
        ):
            self.register_buffer(
                buffer_name,
                torch.zeros(self.num_arms, dtype=torch.float32),
                persistent=True,
            )

        # --- definitive-knowledge ledger TENSORS (source of truth) ----------
        # int8 polarity (0=none, +1=pos, -1=neg); float32 confidence; long
        # knowledge-id hash tokens (0 = unbound). All default-empty so the
        # legacy byte-identical behavior holds until the operator arms the
        # ledger.
        self.register_buffer(
            "definitive_polarity_t",
            torch.zeros(self.num_arms, dtype=torch.int8),
            persistent=True,
        )
        self.register_buffer(
            "definitive_confidence_t",
            torch.zeros(self.num_arms, dtype=torch.float32),
            persistent=True,
        )
        self.register_buffer(
            "knowledge_id_tokens_t",
            torch.zeros(self.num_arms, dtype=torch.long),
            persistent=True,
        )
        # Reverse-lookup bookkeeping for the derived dict/tuple views + the
        # read-only ledger receipt. NEVER model-facing.
        self._knowledge_token_to_id: dict[int, str] = {}

        # Fold any caller-supplied legacy forms into the tensors. ``__post_init__
        # -style normalization: clamp confidence to [0, 1], drop out-of-range
        # arm keys, drop a width-mismatched binding tuple (legacy behavior).
        if arm_knowledge_ids and len(arm_knowledge_ids) != self.num_arms:
            arm_knowledge_ids = ()
        # Prefer the explicit reverse map when supplied (restore path);
        # otherwise rebuild it from the binding tuple for idempotency.
        if knowledge_to_arm is None:
            if arm_knowledge_ids:
                knowledge_to_arm = {
                    str(kid): int(idx)
                    for idx, kid in enumerate(arm_knowledge_ids)
                    if str(kid).strip()
                }
            else:
                knowledge_to_arm = {}
        # Stamp the binding tensors + reverse-lookup bookkeeping.
        if arm_knowledge_ids or knowledge_to_arm:
            self._set_knowledge_to_arm(knowledge_to_arm)
        # Stamp the polarity / confidence tensors from any supplied dicts.
        if definitive_polarity:
            self._set_definitive_polarity(definitive_polarity)
        if definitive_confidence:
            self._set_definitive_confidence(definitive_confidence)

    def _apply(
        self,
        fn: Callable[[torch.Tensor], torch.Tensor],
        recurse: bool = True,
    ) -> "TensorTraumaState":
        """Move hard-knowledge state without narrowing its FP32 evidence.

        These EMAs and confidence tensors are durable learned knowledge, not
        low-precision activation storage.  Preserve their exact FP32 values
        when a surrounding RBO is cast to BF16 so interrupted training can
        cold-resume without changing either positive or negative evidence.
        """

        from resynthesis.quantile_balancing import (
            _apply_fp32_control_buffer_without_narrowing,
        )

        fp32_control_buffers = {
            name: buffer_t
            for name, buffer_t in self._buffers.items()
            if buffer_t is not None and buffer_t.dtype == torch.float32
        }
        result = cast(
            "TensorTraumaState",
            super()._apply(  # type: ignore[no-untyped-call]
                fn,
                recurse=recurse,
            ),
        )
        for name, buffer_t in fp32_control_buffers.items():
            self._buffers[name] = _apply_fp32_control_buffer_without_narrowing(
                buffer_t,
                fn,
            )
        return result

    def grow_prefix_exact(self, num_arms: int) -> None:
        """Grow every per-arm buffer while preserving the old prefix exactly.

        Page catalogs are append-only.  Growth therefore copies the complete
        learned prefix byte-for-byte and zero-initializes only the new suffix.
        Shrinking would destroy learned identities and is rejected.
        """

        requested = int(num_arms)
        if requested < self.num_arms:
            raise ValueError("trauma state cannot shrink its learned prefix")
        if requested == self.num_arms:
            return
        old_width = self.num_arms
        for name, buffer_t in tuple(self.named_buffers(recurse=False)):
            if buffer_t.ndim == 0:
                continue
            if buffer_t.shape[0] != old_width:
                raise RuntimeError(
                    f"trauma buffer {name} does not share the arm prefix"
                )
            suffix_shape = (requested - old_width, *buffer_t.shape[1:])
            if name == "last_success_step":
                suffix_t = buffer_t.new_full(suffix_shape, -1)
            else:
                suffix_t = buffer_t.new_zeros(suffix_shape)
            setattr(self, name, torch.cat((buffer_t, suffix_t), dim=0))
        self.num_arms = requested

    # --- tensor-source-of-truth mutators (internal) -------------------------

    def _set_knowledge_to_arm(self, mapping: dict[str, int]) -> None:
        """Rebuild ``knowledge_id_tokens_t`` + the reverse-lookup map.

        ``mapping`` is ``{knowledge_id_str: arm_index}``. Each entry is hashed
        to a stable non-zero token and stamped at its arm index; arms absent
        from the map are reset to ``DEFINITIVE_UNBOUND_TOKEN`` (0). The
        reverse-lookup ``_knowledge_token_to_id`` is rebuilt so the derived
        dict/tuple views round-trip exactly.
        """

        tokens = self.knowledge_id_tokens_t.new_zeros(self.num_arms)
        rev: dict[int, str] = {}
        for kid, arm in mapping.items():
            arm_i = int(arm)
            if not (0 <= arm_i < self.num_arms):
                continue
            kid_s = str(kid)
            if not kid_s.strip():
                continue
            tok = _knowledge_id_to_token(kid_s)
            tokens[arm_i] = tok
            rev[tok] = kid_s
        self.knowledge_id_tokens_t = tokens
        self._knowledge_token_to_id = rev

    def _set_definitive_polarity(self, mapping: dict[int, str]) -> None:
        """Rebuild ``definitive_polarity_t`` from a ``{arm: polarity_str}`` map.

        Out-of-range arm keys are dropped; unknown polarity strings map to 0
        (none) so a malformed restore cannot corrupt the bank.
        """

        pol = self.definitive_polarity_t.new_zeros(self.num_arms)
        for arm, polarity in mapping.items():
            arm_i = int(arm)
            if not (0 <= arm_i < self.num_arms):
                continue
            pol[arm_i] = int(
                DEFINITIVE_POLARITY_CODE.get(str(polarity).strip().lower(), 0)
            )
        self.definitive_polarity_t = pol

    def _set_definitive_confidence(self, mapping: dict[int, float]) -> None:
        """Rebuild ``definitive_confidence_t`` from a ``{arm: conf}`` map.

        Confidence is clamped to ``[FLOOR, CEIL]``; out-of-range arms dropped.
        """

        conf = self.definitive_confidence_t.new_zeros(self.num_arms)
        for arm, value in mapping.items():
            arm_i = int(arm)
            if not (0 <= arm_i < self.num_arms):
                continue
            conf[arm_i] = max(
                DEFAULT_DEFINITIVE_CONFIDENCE_FLOOR,
                min(DEFAULT_DEFINITIVE_CONFIDENCE_CEIL, float(value)),
            )
        self.definitive_confidence_t = conf

    # --- derived legacy views (read the tensors) ----------------------------

    @property
    def arm_knowledge_ids(self) -> tuple[str, ...]:
        """Derived per-arm knowledge-id tuple (legacy view; reads the tensors).

        Returns ``()`` when NO arm is bound (the legacy default-empty form so
        an un-armed bank is byte-identical to the old behavior, and the
        snapshot omits the additive key). Once at least one arm is bound,
        returns a ``num_arms``-length tuple where each entry is the bound
        knowledge id string or ``""`` for an unbound arm. Reconstructed from
        ``knowledge_id_tokens_t`` + ``_knowledge_token_to_id`` on every read,
        so it always reflects the current tensor state.
        """

        rev = self._knowledge_token_to_id
        tokens = self.knowledge_id_tokens_t
        if tokens.numel() == 0 or not rev:
            # No bindings recorded -> legacy empty-tuple form.
            return ()
        out: list[str] = []
        for tok in tokens.detach().cpu().tolist():
            t = int(tok)
            out.append(rev.get(t, "") if t != DEFINITIVE_UNBOUND_TOKEN else "")
        return tuple(out)

    @arm_knowledge_ids.setter
    def arm_knowledge_ids(self, value: tuple[str, ...]) -> None:
        """Legacy assignment path — rebuilds the binding tensors.

        Kept so ``trauma_state_restore`` and any legacy caller that assigns the
        tuple form still work byte-identically. ``""`` entries are treated as
        unbound (legacy semantics).
        """

        ids = tuple(str(k) for k in (value or ()))
        if ids and len(ids) != self.num_arms:
            # Width mismatch -> drop (legacy behavior, never raise).
            ids = ()
        mapping: dict[str, int] = {
            kid: idx for idx, kid in enumerate(ids) if kid.strip()
        }
        self._set_knowledge_to_arm(mapping)

    @property
    def knowledge_to_arm(self) -> dict[str, int]:
        """Derived ``{knowledge_id: arm_index}`` reverse map (legacy view)."""

        rev = self._knowledge_token_to_id
        tokens = self.knowledge_id_tokens_t
        out: dict[str, int] = {}
        if tokens.numel() == 0:
            return out
        for arm, tok in enumerate(tokens.detach().cpu().tolist()):
            t = int(tok)
            if t == DEFINITIVE_UNBOUND_TOKEN:
                continue
            kid = rev.get(t)
            if kid is not None:
                out[kid] = int(arm)
        return out

    @knowledge_to_arm.setter
    def knowledge_to_arm(self, value: dict[str, int]) -> None:
        """Legacy assignment path — rebuilds the binding tensors."""

        self._set_knowledge_to_arm(dict(value or {}))

    @property
    def definitive_polarity(self) -> dict[int, str]:
        """Derived ``{arm: polarity_str}`` map (legacy view; reads the tensor).

        Only arms whose polarity code is non-zero (positive/negative) appear;
        ``none`` arms are omitted (matches the legacy dict semantics where
        ``mark_definitive`` is the only stamp path).
        """

        pol_t = self.definitive_polarity_t
        out: dict[int, str] = {}
        if pol_t.numel() == 0:
            return out
        for arm, code in enumerate(pol_t.detach().cpu().tolist()):
            c = int(code)
            if c == 0:
                continue
            name = DEFINITIVE_POLARITY_DECODE.get(c)
            if name is not None:
                out[int(arm)] = name
        return out

    @definitive_polarity.setter
    def definitive_polarity(self, value: dict[int, str]) -> None:
        """Legacy assignment path — rebuilds the polarity tensor."""

        self._set_definitive_polarity(dict(value or {}))

    @property
    def definitive_confidence(self) -> dict[int, float]:
        """Derived ``{arm: confidence}`` map (legacy view; reads the tensor).

        Key membership tracks ``definitive_polarity_t`` (a stamped arm always
        has both a polarity and a confidence), so only arms with a non-zero
        polarity code appear here.
        """

        pol_t = self.definitive_polarity_t
        conf_t = self.definitive_confidence_t
        out: dict[int, float] = {}
        if pol_t.numel() == 0 or conf_t.numel() == 0:
            return out
        pol_list = pol_t.detach().cpu().tolist()
        conf_list = conf_t.detach().cpu().tolist()
        n = min(len(pol_list), len(conf_list))
        for arm in range(n):
            if int(pol_list[arm]) == 0:
                continue
            out[int(arm)] = float(conf_list[arm])
        return out

    @definitive_confidence.setter
    def definitive_confidence(self, value: dict[int, float]) -> None:
        """Legacy assignment path — rebuilds the confidence tensor."""

        self._set_definitive_confidence(dict(value or {}))



def _coerce_arm_index(arm_index_t: Tensor, num_arms: int) -> Tensor:
    """Return only exact in-catalog arm positions.

    An invalid catalog identity is evidence about no arm.  Clamping or modulo
    would fabricate a successful/failing observation for an unrelated page,
    so out-of-range identities are rejected rather than rewritten.
    """
    idx = arm_index_t.detach().reshape(-1).to(dtype=torch.long)
    return idx[idx.ge(0) & idx.lt(num_arms)]


def _coerce_magnitude(
    magnitude_t: Tensor | float, ref: Tensor
) -> Tensor:
    """Normalize a magnitude argument to a scalar tensor matching ``ref``."""

    if isinstance(magnitude_t, Tensor):
        return magnitude_t.reshape(()).to(
            device=ref.device, dtype=ref.dtype
        )
    return ref.new_tensor(float(magnitude_t))


def update_trauma_from_verified_outcome(
    state: TensorTraumaState,
    *,
    arm_index_t: Tensor,
    frontier_weight_t: Tensor,
    pro_t: Tensor,
    anti_t: Tensor,
    evidence_confidence_t: Tensor,
    knowledge_axis_t: Tensor,
    behavior_axis_t: Tensor,
) -> Tensor:
    """Commit verified, frontier-weighted evidence for *future* routes.

    This boundary is called only after the target-independent forward and its
    KLA/behavior verifier have completed.  It never returns route logits for
    that completed forward; it mutates persistent buffers consumed by the next
    call.  Every selected catalog position receives its exact normalized
    frontier mass.  Invalid positions, non-finite weights, and zero-mass rows
    are rejected rather than clamped onto a different page.

    Returns the unique catalog positions that accepted evidence.  The return is
    tensor-native so callers can record invalid/empty evidence without a host
    identity rewrite.
    """

    if arm_index_t.ndim != 1 or frontier_weight_t.ndim != 1:
        raise ValueError("verified trauma outcome requires one-dimensional arms/weights")
    if arm_index_t.numel() != frontier_weight_t.numel():
        raise ValueError("verified trauma outcome arm/weight geometry differs")
    device = state.fail_ema.device
    raw_index_t = arm_index_t.detach().to(device=device, dtype=torch.long)
    raw_weight_t = frontier_weight_t.detach().to(
        device=device,
        dtype=state.fail_ema.dtype,
    )
    valid_t = (
        raw_index_t.ge(0)
        & raw_index_t.lt(state.num_arms)
        & torch.isfinite(raw_weight_t)
        & raw_weight_t.gt(0)
    )
    valid_index_t = raw_index_t[valid_t]
    valid_weight_t = raw_weight_t[valid_t]
    if valid_index_t.numel() == 0:
        return raw_index_t.new_empty(0)

    unique_result = cast(
        tuple[Tensor, Tensor],
        torch.unique(
            valid_index_t,
            sorted=True,
            return_inverse=True,
        ),
    )
    index_t, inverse_t = unique_result
    mass_t = state.fail_ema.new_zeros(index_t.numel())
    mass_t.scatter_add_(0, inverse_t, valid_weight_t)

    def _verified_unit(value_t: Tensor) -> Tensor:
        value = value_t.detach().reshape(()).to(
            device=device,
            dtype=state.fail_ema.dtype,
        )
        return torch.nan_to_num(value, nan=0.0, posinf=1.0, neginf=0.0).clamp(
            0.0,
            1.0,
        )

    pro_unit_t = _verified_unit(pro_t)
    anti_unit_t = _verified_unit(anti_t)
    confidence_t = _verified_unit(evidence_confidence_t)
    knowledge_t = _verified_unit(knowledge_axis_t)
    behavior_t = _verified_unit(behavior_axis_t)
    pro_mass_t = mass_t * pro_unit_t * confidence_t
    anti_mass_t = mass_t * anti_unit_t * confidence_t

    with torch.no_grad():
        state.global_step.add_(1)
        step_now_t = state.global_step.reshape(()).to(dtype=torch.long)
        prior_fail_t = state.fail_ema.index_select(0, index_t)
        prior_success_t = state.success_ema.index_select(0, index_t)
        next_fail_t = prior_fail_t * state.fail_decay + anti_mass_t
        next_success_t = prior_success_t * state.success_decay + pro_mass_t
        state.fail_ema.index_copy_(0, index_t, next_fail_t)
        state.success_ema.index_copy_(0, index_t, next_success_t)

        prior_peak_t = state.peak_fail_ema.index_select(0, index_t)
        struggled_t = prior_peak_t.ge(
            state.peak_fail_ema.new_tensor(
                DEFAULT_HARD_WON_STRUGGLE_THRESHOLD
            )
        )
        next_peak_t = torch.maximum(prior_peak_t, next_fail_t)
        next_peak_t = torch.where(
            pro_mass_t.gt(0),
            next_peak_t * (1.0 - pro_unit_t),
            next_peak_t,
        )
        state.peak_fail_ema.index_copy_(0, index_t, next_peak_t)
        hard_won_add_t = pro_mass_t * struggled_t.to(dtype=pro_mass_t.dtype)
        state.hard_won_ema.index_add_(0, index_t, hard_won_add_t)
        prior_success_step_t = state.last_success_step.index_select(0, index_t)
        next_success_step_t = torch.where(
            pro_mass_t.gt(0),
            step_now_t.expand_as(prior_success_step_t),
            prior_success_step_t,
        )
        state.last_success_step.index_copy_(
            0,
            index_t,
            next_success_step_t,
        )

        for name, delta_t in (
            ("knowledge_pro_ema", pro_mass_t * knowledge_t),
            ("knowledge_anti_ema", anti_mass_t * knowledge_t),
            ("behavior_pro_ema", pro_mass_t * behavior_t),
            ("behavior_anti_ema", anti_mass_t * behavior_t),
            ("verified_exposure_ema", mass_t),
            ("evidence_confidence_ema", mass_t * confidence_t),
        ):
            buffer_t = getattr(state, name)
            previous_t = buffer_t.index_select(0, index_t)
            buffer_t.index_copy_(
                0,
                index_t,
                previous_t * state.success_decay + delta_t,
            )
    return index_t


def update_trauma_from_outcome(
    state: TensorTraumaState,
    *,
    arm_index_t: Tensor,
    success_t: Tensor | bool | float,
    magnitude_t: Tensor | float = 1.0,
) -> None:
    """Fold one outcome into the trauma bank — the learn_loop boundary call.

    WHAT:
        Increments ``fail_ema`` (negative trauma) on a failed outcome and
        ``success_ema`` (positive trauma) on a successful outcome, with the
        configured per-arm exponential decay applied first. Advances the
        global step counter and stamps the per-arm last-success step.

    WHY:
        This is the single ingest surface the training loop calls at the
        outcome boundary (inside the existing fail-open retention try/except).
        Keeping it to one call keeps the wiring trivial and the cost O(1)
        per arm updated.

    HOW (adapted from trauma_informed_trainer + anti_system_metrics):
      * Negative path: ``fail_ema <- fail_decay * fail_ema + magnitude`` for
        each failed arm. This is the anti-Thompson-style repeated-failure
        accumulation, with decay so trauma fades when an arm recovers.
      * Positive path: ``success_ema <- success_decay * success_ema +
        magnitude`` for each successful arm, AND the global step is stamped
        into ``last_success_step`` so the cooldown gate can read it.
      * Both paths are tensor-native ``no_grad`` index ops; no Python loop
        over arms, no host roundtrip.

    Invalid arm identities are dropped.  This compatibility surface treats its
    caller as a verified knowledge outcome; production behavior/KLA callers use
    ``update_trauma_from_verified_outcome`` to retain their separate axes.

    Args:
      state:        the ``TensorTraumaState`` bank to mutate in place.
      arm_index_t:  1-D long tensor of arm indices that produced this outcome.
      success_t:    per-call success flag/value (tensor / bool / float). The
                    SAME success value is applied to every supplied arm index
                    (call once per distinct outcome class if you need mixed).
      magnitude_t:  how much this outcome weighs (default 1.0). Pass a tensor
                    to make it differentiable for offline analysis; the live
                    learn_loop path passes a detached scalar.
    """

    if not isinstance(arm_index_t, Tensor):
        return
    idx = _coerce_arm_index(
        arm_index_t.to(device=state.fail_ema.device),
        state.num_arms,
    )
    if idx.numel() == 0:
        return
    mag = torch.nan_to_num(
        _coerce_magnitude(magnitude_t, state.fail_ema),
        nan=0.0,
        posinf=0.0,
        neginf=0.0,
    ).clamp_min(0.0)
    success_value_t = (
        success_t.detach().reshape(()).to(
            device=state.fail_ema.device,
            dtype=state.fail_ema.dtype,
        )
        if isinstance(success_t, Tensor)
        else state.fail_ema.new_tensor(float(success_t))
    )
    success_value_t = torch.nan_to_num(
        success_value_t,
        nan=0.0,
        posinf=1.0,
        neginf=0.0,
    ).clamp(0.0, 1.0)
    update_trauma_from_verified_outcome(
        state,
        arm_index_t=idx,
        frontier_weight_t=mag.expand(idx.numel()),
        pro_t=success_value_t,
        anti_t=1.0 - success_value_t,
        evidence_confidence_t=state.fail_ema.new_ones(()),
        knowledge_axis_t=state.fail_ema.new_ones(()),
        behavior_axis_t=state.fail_ema.new_zeros(()),
    )


def trauma_pressure_t(
    state: TensorTraumaState,
    *,
    scale: float = DEFAULT_TRAUMA_PRESSURE_SCALE,
) -> Tensor:
    """Negative-trauma repulsion bias for the quantile router.

    Returns a per-arm non-negative bias increment (same sign convention as
    ``TensorAntiThompsonRegistry.anti_bias_for_quantile_router``: the router
    ADDS this to ``expert_bias_t`` to push AWAY from traumatized arms).

    Adapted from ``trauma_scaler.forward``: the source forms a positive
    multiplier from a centered trauma signal via ``softplus``. Here we reuse
    the centered+softplus idea but invert it into a router-pushing bias:
    arms with high ``fail_ema`` get a large positive increment (repelled),
    arms with no failures get zero. The ``softplus`` keeps it smooth and
    positive (no sign flips that could attract a traumatized arm by mistake).

    Args:
      state: the trauma bank to read.
      scale: maximum bias magnitude (matches ``DEFAULT_ANTI_BIAS_SCALE`` family).

    Returns:
      ``[num_arms]`` float32 tensor of non-negative bias increments.
    """

    fails = state.fail_ema.clamp_min(0.0)
    if fails.numel() == 0 or float(fails.max().item()) <= 0.0:
        return fails  # all-zero shortcut; preserves device/dtype
    # Center on the max so the MOST traumatized arm gets the full scale and
    # less-traumatized arms get proportionally less. softplus(smooth) keeps it
    # differentiable for offline analysis and strictly non-negative. Mask by
    # the raw fail EMA so never-failed arms get EXACTLY zero repulsion
    # (softplus(0) has a ~0.693 floor that would otherwise leak a tiny bias to
    # clean arms; the mask makes the anti-collapse boundary exact).
    ever_failed = fails > 0.0
    centered = fails - fails.max().detach()
    pressure = F.softplus(centered * 4.0) * ever_failed.to(dtype=fails.dtype)
    peak = pressure.max().clamp_min(DEFAULT_TRAUMA_EPS)
    return (pressure / peak) * float(scale)


def positive_reinforcement_t(
    state: TensorTraumaState,
    *,
    scale: float = DEFAULT_TRAUMA_POSITIVE_SCALE,
    cooldown_steps: int = DEFAULT_TRAUMA_COOLDOWN_STEPS,
    high_fail_threshold: float = DEFAULT_TRAUMA_HIGH_FAIL_THRESHOLD,
    max_fraction: float = DEFAULT_TRAUMA_MAX_FRACTION,
) -> Tensor:
    """Positive-trauma reinforcement bias for the quantile router (tempered).

    Returns a per-arm non-negative bias increment that ATTRACTS the router
    toward recently-successful arms. This is the "positive trauma" half of
    the operator's requirement.

    THREE anti-collapse guards (the "tempered" part), each adapted from a
    source idea:

      1. COOLDOWN GATE (from trauma_informed_trainer.adaptation_cooldown):
         If ANY arm is still actively traumatized (``fail_ema`` above
         ``high_fail_threshold``) OR any successful arm succeeded fewer than
         ``cooldown_steps`` ago, return all-zeros. Positive reinforcement
         only fires once the system has genuinely cleared its failures.

      2. MAX-FRACTION CAP (from anti_system_metrics diversity/anti-collapse):
         No single arm may capture more than ``max_fraction`` of the total
         reinforcement mass. Surplus is redistributed uniformly so a dominant
         arm cannot starve the others -> prevents collapse to one arm.

      3. SUCCESS-WEIGHTED (from trauma_scaler softplus): raw signal is the
         success EMA passed through softplus so the bias is smooth, positive,
         and differentiable; recent repeated successes dominate.

    Args:
      state:               the trauma bank to read.
      scale:               maximum total reinforcement mass.
      cooldown_steps:      min steps since last success before reinforcing.
      high_fail_threshold: fail_ema above which positive reinforcement is
                           globally suppressed (system still traumatized).
      max_fraction:        cap on any one arm's share of reinforcement mass.

    Returns:
      ``[num_arms]`` float32 tensor of non-negative bias increments.
    """

    successes = state.success_ema.clamp_min(0.0)
    if successes.numel() == 0 or float(successes.max().item()) <= 0.0:
        return successes  # nothing to reinforce; preserves device/dtype

    # Guard 1: cooldown gate. If the system is still traumatized, or no arm
    # has been stable-successful for long enough, do NOT reinforce.
    step_now = state.global_step.detach().reshape(()).to(dtype=torch.long)
    still_traumatized = bool(
        state.fail_ema.clamp_min(0.0).max().item()
        >= float(high_fail_threshold)
    )
    # last_success_step is -1 for arms that never succeeded; ignore those.
    succeeded_mask = state.last_success_step >= 0
    if bool(succeeded_mask.any().item()):
        steps_since = step_now - state.last_success_step.clamp_min(0)
        # Only consider arms that actually succeeded.
        steps_since = steps_since * succeeded_mask.to(dtype=steps_since.dtype)
        min_steps_since = float(steps_since[succeeded_mask].min().item())
    else:
        min_steps_since = float("inf")
    if still_traumatized or min_steps_since < float(cooldown_steps):
        return successes.new_zeros(successes.shape)

    # Guard 3: smooth, positive, success-weighted signal. Only arms that have
    # EVER succeeded receive any reinforcement mass — softplus(0) has a
    # positive floor (~0.693) that would otherwise leak reinforcement to arms
    # with no success history. Mask by the raw success EMA BEFORE softplus so
    # never-successful arms are exactly zero (clean anti-collapse boundary).
    ever_succeeded = successes > 0.0
    raw = F.softplus(successes * 4.0) * ever_succeeded.to(dtype=successes.dtype)
    total = raw.sum().clamp_min(DEFAULT_TRAUMA_EPS)
    shares = raw / total

    # Guard 2: max-fraction cap with surplus redistribution (water-filling).
    # No single arm may capture more than ``max_fraction`` of the mass. Surplus
    # stripped from over-cap arms is redistributed to OTHER arms that have
    # signal (``raw > 0``), proportionally to their existing share — never to
    # arms with zero signal (which must stay at exactly zero so the
    # anti-collapse boundary stays exact). We iterate a few times because
    # redistribution can push an under-cap arm over the cap. This is the
    # tensor-native, branch-free water-filling loop; bounded by a fixed count
    # so it can never spin. The cap is a SOFT anti-collapse guard: if only one
    # arm has signal, the cap does NOT force its mass to vanish — it just
    # limits how dominant it can be relative to the runner-up.
    cap = float(max_fraction)
    if cap <= 0.0:
        return successes.new_zeros(successes.shape)
    if cap < 1.0:
        signal_mask = raw > 0.0
        signal_count = signal_mask.to(dtype=shares.dtype).sum()
        # If only one arm has signal, the cap is meaningless (no one to
        # redistribute to); skip capping entirely so we never waste mass.
        if float(signal_count.item()) > 1.0:
            for _ in range(8):  # bounded water-filling; converges fast here
                over_mask = shares > cap
                if not bool(over_mask.any().item()):
                    break
                excess = (
                    torch.where(
                        over_mask, shares - cap, shares.new_zeros(())
                    )
                ).sum()
                shares = torch.minimum(
                    shares, shares.new_full((), cap)
                )
                # Redistribute the excess ONLY to other signal arms still under
                # cap, proportional to their current share (so a stronger
                # runner-up absorbs more of the stripped mass).
                under_signal = signal_mask & (shares < cap)
                weights = shares * under_signal.to(dtype=shares.dtype)
                w_sum = weights.sum().clamp_min(DEFAULT_TRAUMA_EPS)
                shares = shares + excess * (weights / w_sum)
            shares = torch.minimum(
                shares, shares.new_full((), cap)
            )

    return shares * float(scale)


def trauma_bias_t(
    state: TensorTraumaState,
    *,
    pressure_scale: float = DEFAULT_TRAUMA_PRESSURE_SCALE,
    positive_scale: float = DEFAULT_TRAUMA_POSITIVE_SCALE,
    cooldown_steps: int = DEFAULT_TRAUMA_COOLDOWN_STEPS,
    high_fail_threshold: float = DEFAULT_TRAUMA_HIGH_FAIL_THRESHOLD,
    max_fraction: float = DEFAULT_TRAUMA_MAX_FRACTION,
) -> Tensor:
    """Signed TRAUMA bias in quantile-beta coordinates.

    ``QuantileBalancingRouter.biased_scores`` computes ``scores - beta``.
    Negative definitive knowledge therefore contributes positive beta
    (repulsion), while positive definitive knowledge contributes negative beta
    (attraction).  Keeping those poles separate prevents the former
    ``pressure + positive`` sign error that suppressed hard-won knowledge.

    FAIL-OPEN by construction: both sub-functions are pure tensor ops.
    """

    pressure = trauma_pressure_t(state, scale=pressure_scale)
    positive = positive_reinforcement_t(
        state,
        scale=positive_scale,
        cooldown_steps=cooldown_steps,
        high_fail_threshold=high_fail_threshold,
        max_fraction=max_fraction,
    )
    return pressure - positive


def apply_trauma_bias_to_quantile_router(
    router: "torch.nn.Module",
    state: TensorTraumaState,
    *,
    pressure_scale: float = DEFAULT_TRAUMA_PRESSURE_SCALE,
    positive_scale: float = DEFAULT_TRAUMA_POSITIVE_SCALE,
    cooldown_steps: int = DEFAULT_TRAUMA_COOLDOWN_STEPS,
    high_fail_threshold: float = DEFAULT_TRAUMA_HIGH_FAIL_THRESHOLD,
    max_fraction: float = DEFAULT_TRAUMA_MAX_FRACTION,
) -> Tensor:
    """Nudge ``expert_bias_t`` using combined trauma (in-place), bridge idiom.

    Mirrors ``apply_anti_bias_to_quantile_router`` exactly in shape and
    fail-open behavior. The router must expose ``expert_bias_t`` with width
    matching ``state.num_arms``. Returns a detached clone of the updated bias.
    """

    bias_delta = trauma_bias_t(
        state,
        pressure_scale=pressure_scale,
        positive_scale=positive_scale,
        cooldown_steps=cooldown_steps,
        high_fail_threshold=high_fail_threshold,
        max_fraction=max_fraction,
    )
    expert_bias_t = getattr(router, "expert_bias_t", None)
    if not isinstance(expert_bias_t, Tensor):
        raise TypeError("quantile router has no tensor expert bias")
    if bias_delta.numel() != expert_bias_t.numel():
        raise ValueError("trauma state width differs from quantile router")
    with torch.no_grad():
        expert_bias_t.add_(bias_delta.to(device=expert_bias_t.device))
    return expert_bias_t.detach().clone()


def negative_hard_knowledge_mask_t(
    state: TensorTraumaState,
    *,
    threshold: float = DEFAULT_HARD_KNOWLEDGE_FAIL_THRESHOLD,
) -> Tensor:
    """Boolean mask — arms carrying **negative hard knowledge** (definitive gap)."""

    from resynthesis.hard_knowledge_surface import hard_knowledge_surface_packet_t

    packet = hard_knowledge_surface_packet_t(
        state, negative_threshold=threshold, top_k=0
    )
    return packet.negative_gap_mask_t.gt(0.0)


def positive_hard_knowledge_mask_t(
    state: TensorTraumaState,
    *,
    threshold: float = DEFAULT_HARD_WON_EMA_THRESHOLD,
) -> Tensor:
    """Boolean mask — arms carrying **positive hard knowledge** (hard-won win)."""

    from resynthesis.hard_knowledge_surface import hard_knowledge_surface_packet_t

    packet = hard_knowledge_surface_packet_t(
        state, positive_threshold=threshold, top_k=0
    )
    return packet.positive_hard_won_mask_t.gt(0.0)


def negative_hard_knowledge_arm_indices_t(
    state: TensorTraumaState,
    *,
    threshold: float = DEFAULT_HARD_KNOWLEDGE_FAIL_THRESHOLD,
    top_k: int = 64,
) -> Tensor:
    """Top arms by negative hard-knowledge pressure (gaps to learn)."""

    from resynthesis.hard_knowledge_surface import hard_knowledge_surface_packet_t

    return hard_knowledge_surface_packet_t(
        state, negative_threshold=threshold, top_k=top_k
    ).negative_arm_indices_t


def positive_hard_knowledge_arm_indices_t(
    state: TensorTraumaState,
    *,
    threshold: float = DEFAULT_HARD_WON_EMA_THRESHOLD,
    top_k: int = 64,
) -> Tensor:
    """Top arms by positive hard-knowledge mass (definitive earned capability)."""

    from resynthesis.hard_knowledge_surface import hard_knowledge_surface_packet_t

    return hard_knowledge_surface_packet_t(
        state, positive_threshold=threshold, top_k=top_k
    ).positive_arm_indices_t


def definitive_hard_knowledge_surface_t(
    state: TensorTraumaState,
    *,
    negative_scale: float = DEFAULT_TRAUMA_PRESSURE_SCALE,
    positive_scale: float = DEFAULT_TRAUMA_POSITIVE_SCALE,
) -> Tensor:
    """Combined hard-knowledge routing surface (tensor-native)."""

    from resynthesis.hard_knowledge_surface import hard_knowledge_surface_packet_t

    return hard_knowledge_surface_packet_t(
        state,
        negative_scale=negative_scale,
        positive_scale=positive_scale,
    ).definitive_surface_t


def hard_knowledge_surface_receipt(
    state: TensorTraumaState,
    *,
    loop_id: str = "",
    iteration: int = 0,
    top_k: int = 32,
) -> dict[str, object]:
    """Boundary receipt — JSON adapter only (hot path uses ``HardKnowledgeSurfacePacket``)."""

    from resynthesis.hard_knowledge_surface import hard_knowledge_surface_packet_t

    receipt: dict[str, object] = {
        "schema": HARD_KNOWLEDGE_RECEIPT_SCHEMA,
        "traumaSchema": TRAUMA_SYSTEM_SCHEMA,
        "loopId": str(loop_id),
        "iteration": int(iteration),
    }
    try:
        packet = hard_knowledge_surface_packet_t(state, top_k=top_k)
        receipt["numArms"] = int(state.num_arms)
        receipt["negativeHardKnowledgeCount"] = int(packet.negative_count_t.reshape(()).item())
        receipt["positiveHardKnowledgeCount"] = int(packet.positive_count_t.reshape(()).item())
        receipt["negativeHardKnowledgeArmIds"] = packet.negative_arm_indices_t.tolist()
        receipt["positiveHardKnowledgeArmIds"] = packet.positive_arm_indices_t.tolist()
        receipt["definitiveGapArms"] = receipt["negativeHardKnowledgeArmIds"]
        receipt["hardWonArms"] = receipt["positiveHardKnowledgeArmIds"]
        receipt["tensorSchema"] = packet.schema
    except Exception as receipt_error:  # pragma: no cover - fail-open
        receipt["receiptError"] = repr(receipt_error)
    return receipt


def trauma_state_to_receipt(
    state: TensorTraumaState,
    *,
    loop_id: str = "",
    iteration: int = 0,
) -> dict[str, object]:
    """Schema-sealed boundary receipt for telemetry / persistence.

    Adapted from ``loss_telemetry``'s fail-open receipt pattern: every value
    is a plain Python scalar extracted under a try/except so a malformed
    tensor can never break JSON serialization. This is the receipt the
    learn_loop boundary appends to the trauma ledger.

    Returns a dict with the schema constant and a compact, host-safe summary
    (per-arm max/mean trauma, total successes, active-trauma arm count). It
    does NOT duplicate the full per-arm tensors — only scalars — so the
    ledger stays small.
    """

    receipt: dict[str, object] = {
        "schema": TRAUMA_SYSTEM_SCHEMA,
        "loopId": str(loop_id),
        "iteration": int(iteration),
    }
    try:
        receipt["numArms"] = int(state.num_arms)
        receipt["globalStep"] = int(state.global_step.item())
        receipt["failEmaMax"] = float(state.fail_ema.max().item())
        receipt["failEmaMean"] = float(state.fail_ema.mean().item())
        receipt["successEmaMax"] = float(state.success_ema.max().item())
        receipt["successEmaMean"] = float(state.success_ema.mean().item())
        receipt["activelyTraumatizedArms"] = int(
            (state.fail_ema >= DEFAULT_TRAUMA_HIGH_FAIL_THRESHOLD).sum().item()
        )
        receipt["armsWithSuccess"] = int(
            (state.last_success_step >= 0).sum().item()
        )
        hk = hard_knowledge_surface_receipt(
            state, loop_id=loop_id, iteration=iteration, top_k=16
        )
        receipt["hardKnowledge"] = {
            "negativeCount": hk.get("negativeHardKnowledgeCount"),
            "positiveCount": hk.get("positiveHardKnowledgeCount"),
            "definitiveGapArms": hk.get("negativeHardKnowledgeArmIds"),
            "hardWonArms": hk.get("positiveHardKnowledgeArmIds"),
        }
    except Exception as receipt_error:  # pragma: no cover - fail-open
        receipt["receiptError"] = repr(receipt_error)
    return receipt


# ============================================================================
# ADDITIVE DEFINITIVE-KNOWLEDGE LEDGER (off-by-default, fail-open everywhere)
# ----------------------------------------------------------------------------
# These functions are NEW and ADDITIVE. None of them touch the existing trauma
# tensors or routing math; they layer a knowledge-identity + confidence +
# preservation + must-learn-priority + ledger view on top of the same per-arm
# state. The hot path (update_trauma_from_outcome, trauma_pressure_t,
# positive_reinforcement_t, trauma_bias_t, the router fold) is byte-identical
# when these are unused. The whole layer is gated off by
# ``NNF_TRAUMA_DEFINITIVE_LEDGER=1`` and every public function FAILS-OPEN to an
# empty/clean result when the gate is unset or any error occurs.
# ============================================================================


def _definitive_ledger_enabled_boundary() -> bool:
    """Read the off-by-default env gate for the definitive-knowledge layer.

    Fail-open: any read error returns False (layer off) so a malformed env can
    never break the caller. The gate is OFF by default; the operator arms it
    when they want to persist/consume knowledge bindings + definitive stamps.
    """

    try:
        import os

        return str(os.environ.get(ENV_TRAUMA_DEFINITIVE_LEDGER, "")).strip() in {
            "1",
            "true",
            "True",
            "TRUE",
        }
    except Exception:  # pragma: no cover - fail-open
        return False


def _validate_arm_index(state: TensorTraumaState, arm_index: int) -> int:
    """Validate an arm index against the trauma bank width, fail-open.

    Returns the clamped index in [0, num_arms). A negative or oversized index
    is rejected by raising IndexError so the caller's try/except can fail-open.
    Kept separate from ``_coerce_arm_index`` (which clamps tensor inputs) so
    the ledger API stays plain-Python and never silently rewrites a bad id.
    """

    if not isinstance(arm_index, (int,)) or arm_index < 0 or arm_index >= state.num_arms:
        raise IndexError(
            f"arm_index {arm_index} out of range [0, {state.num_arms})"
        )
    return int(arm_index)


def bind_knowledge_to_arm(
    state: TensorTraumaState, arm_index: int, knowledge_id: str
) -> None:
    """Bind one arm to a knowledge identity (additive, fail-open).

    WHAT / WHY:
        Trauma arms must answer "which hard KNOWLEDGE is positive-definitive /
        negative-definitive", not "which abstract index failed". This records
        the knowledge identity (e.g. a capability key / page knowledge hash /
        CWE label) for ``arm_index`` so the ledger receipt and must-learn
        priority surface knowledge ids, not bare integers.

    HOW:
        Rebuilds the ``arm_knowledge_ids`` tuple and the ``knowledge_to_arm``
        reverse map. If ``knowledge_id`` is empty/blank the binding is dropped
        (legacy behavior for that arm). Re-binding an existing knowledge id to
        a NEW arm silently migrates it (last writer wins) so a re-keyed
        capability cannot ghost in two places.

    FAIL-OPEN: any error returns without mutating state. No-op when the
    definitive-ledger gate is OFF (binding is harmless but we keep the layer
    inert until armed, matching the off-by-default invariant).
    """

    try:
        if not _definitive_ledger_enabled_boundary():
            return
        arm = _validate_arm_index(state, arm_index)
        kid = str(knowledge_id).strip()
        if not kid:
            return
        # Build the new reverse map. Migrate first so a re-keyed capability
        # cannot appear under two arms (last writer wins).
        new_map = dict(state.knowledge_to_arm)
        for prior_kid, prior_arm in list(new_map.items()):
            if prior_kid == kid and prior_arm != arm:
                new_map.pop(prior_kid, None)
        new_map[kid] = arm
        # Write through the tensor-source-of-truth mutator (rebuilds
        # ``knowledge_id_tokens_t`` + ``_knowledge_token_to_id``). The derived
        # ``arm_knowledge_ids`` tuple + ``knowledge_to_arm`` dict views reflect
        # this automatically.
        state._set_knowledge_to_arm(new_map)
    except Exception:  # pragma: no cover - fail-open
        return


def arm_for_knowledge(
    state: TensorTraumaState, knowledge_id: str
) -> int | None:
    """Reverse-lookup: which arm carries ``knowledge_id`` (None if unbound).

    Fail-open: returns None on any error or unknown id. Pure read; never
    mutates state. Works regardless of the ledger gate (a read of an empty
    map is a clean None), so callers can probe bindings defensively.
    """

    try:
        kid = str(knowledge_id).strip()
        if not kid:
            return None
        found = state.knowledge_to_arm.get(kid)
        return int(found) if found is not None else None
    except Exception:  # pragma: no cover - fail-open
        return None


def mark_definitive(
    state: TensorTraumaState,
    arm_index: int,
    polarity: str,
    confidence: float,
    *,
    confidence_floor: float = DEFAULT_DEFINITIVE_CONFIDENCE,
) -> bool:
    """Stamp one arm ``definitive`` with a polarity + confidence (anchor class).

    WHAT / WHY (doctrine):
        Only HIGH-confidence, verification-survived knowledge becomes
        "definitive" — the anchor class. This is the ONLY stamp path. The
        caller MUST be a verification boundary (e.g. a probe pass AFTER
        repeated struggle -> positive-definitive; a verified reproduction of
        a failure -> negative-definitive). It is NEVER auto-stamped on noise.

    HOW:
      * ``polarity`` must be ``DEFINITIVE_POLARITY_POSITIVE`` ("positive" =
        hard-won truth to PRESERVE) or ``DEFINITIVE_POLARITY_NEGATIVE``
        ("negative" = hard anti-pattern to AVOID). Any other value is rejected
        (returns False) so the stamp cannot corrupt the ledger.
      * ``confidence`` is clamped to [0, 1]. If the clamped confidence is below
        ``confidence_floor`` (default ``DEFAULT_DEFINITIVE_CONFIDENCE``) the
        stamp is REJECTED (returns False) — the arm stays a candidate, not an
        anchor. This is the doctrine gate: definitive = high-confidence.
      * On success, stamps ``definitive_polarity[arm]`` and
        ``definitive_confidence[arm]`` and returns True.

    FAIL-OPEN: returns False on any error. No-op (returns False) when the
    definitive-ledger gate is OFF, so a verification boundary cannot write the
    ledger before the operator arms it.

    Returns:
      True if the stamp landed; False if rejected (low confidence / bad
      polarity / gate off / error). Callers should treat False as "stay a
      candidate, do not promote to anchor".
    """

    try:
        if not _definitive_ledger_enabled_boundary():
            return False
        arm = _validate_arm_index(state, arm_index)
        pol = str(polarity).strip().lower()
        if pol not in {DEFINITIVE_POLARITY_POSITIVE, DEFINITIVE_POLARITY_NEGATIVE}:
            return False
        conf = max(
            DEFAULT_DEFINITIVE_CONFIDENCE_FLOOR,
            min(DEFAULT_DEFINITIVE_CONFIDENCE_CEIL, float(confidence)),
        )
        if not (conf + DEFAULT_MUST_LEARN_COVERAGE_EPS >= float(confidence_floor)):
            return False
        # Write the TENSOR source of truth directly (int8 polarity code +
        # float confidence at this arm). The derived ``definitive_polarity`` /
        # ``definitive_confidence`` dict views reflect this automatically.
        state.definitive_polarity_t[arm] = int(
            DEFINITIVE_POLARITY_CODE.get(pol, 0)
        )
        state.definitive_confidence_t[arm] = float(conf)
        return True
    except Exception:  # pragma: no cover - fail-open
        return False


def positive_definitive_preservation_weights(
    state: TensorTraumaState,
    *,
    scale: float = DEFAULT_POSITIVE_DEFINITIVE_PRESERVATION_SCALE,
) -> Tensor:
    """Per-arm weights for hard-won POSITIVE-definitive knowledge (anti-forgetting).

    WHAT / WHY (doctrine):
        POSITIVE trauma = PRESERVE. Hard-won positive-definitive knowledge
        (the truths/abilities the model MASTERED through difficulty and that
        survived verification) must NOT be lost during further training. This
        returns the per-arm weight training should use for anti-forgetting /
        replay priority so the model never loses its hard-won truths.

    HOW:
        Weight = hard_won_ema (normalized) * definitive_confidence, restricted
        to arms stamped ``positive``-definitive. Cheap O(num_arms) scalar
        vector. Arms not positive-definitive get EXACTLY zero (clean boundary
        so preservation never leaks to unverified arms).

    FAIL-OPEN: returns a clean zero vector on any error, or when the
    definitive-ledger gate is OFF (preservation is inert until armed). The
    caller is responsible for folding this into its replay/anti-forgetting
    schedule; this function only produces the per-arm scalar surface.

    Args:
      state: the trauma bank to read.
      scale: scalar multiplier on the normalized preservation surface.

    Returns:
      ``[num_arms]`` float32 tensor of non-negative preservation weights.
    """

    try:
        out = state.hard_won_ema.new_zeros(state.num_arms)
        if not _definitive_ledger_enabled_boundary():
            return out
        hard_won = state.hard_won_ema.clamp_min(0.0)
        # Restrict to positive-definitive arms directly from the polarity
        # tensor (int8 code +1 == positive) and read the confidence tensor
        # natively. This is the model-facing tensor path; arms not positive-
        # definitive get EXACTLY zero via the pos_mask multiply.
        pos_mask = (
            state.definitive_polarity_t.to(
                device=hard_won.device, dtype=hard_won.dtype
            )
            == float(DEFINITIVE_POLARITY_CODE[DEFINITIVE_POLARITY_POSITIVE])
        ).to(dtype=hard_won.dtype)
        conf_vec = state.definitive_confidence_t.to(
            device=hard_won.device, dtype=hard_won.dtype
        )
        raw = hard_won * conf_vec * pos_mask
        peak = raw.max()
        if float(peak.item()) > 0.0:
            raw = raw / peak.clamp_min(DEFAULT_TRAUMA_EPS)
        return raw * float(scale)
    except Exception:  # pragma: no cover - fail-open
        try:
            return state.hard_won_ema.new_zeros(state.num_arms)
        except Exception:  # pragma: no cover - fail-open
            return torch.zeros(state.num_arms, dtype=torch.float32)


def hard_knowledge_must_learn_priority(
    state: TensorTraumaState,
    coverage_payload: dict[str, float] | None = None,
    *,
    top_k: int = DEFAULT_MUST_LEARN_PRIORITY_TOP_K,
) -> dict[str, object]:
    """Must-learn curriculum priority for the hard definitive knowledge.

    WHAT / WHY (doctrine):
        The definitive hard knowledge gets TOP learning priority. This ranks
        BOTH polarities — positive-definitive (reinforce/preserve) and
        negative-definitive (contrast/avoid) — and returns a priority map the
        coverage-pressure targeting + the MitM scaffold consume to point onto
        the RIGHT hard targets. This is the curriculum driver.

    HOW (formula, per arm):
        trauma_level    = fail_ema + peak_fail_ema + hard_won_ema
                          (negative gap pressure + positive hard-won mass)
        confidence      = definitive_confidence (1.0 if NOT stamped yet, so
                          un-stamped hard knowledge is still ranked by trauma;
                          definitive arms get their stamped confidence which is
                          >= DEFAULT_DEFINITIVE_CONFIDENCE by construction)
        coverage        = coverage_payload.get(knowledge_id, 0.0) in [0, 1]
                          (1.0 = fully covered; default 0.0 = uncovered)
        priority        = trauma_level * confidence * (1 - coverage)

        The (1 - coverage) term is what makes this a CURRICULUM priority: a
        hard arm that is already covered drops out, steering the learner onto
        the next hard target. Confidence multiplies so verified-definitive
        anchors outrank unverified candidates at equal trauma.

    OUTPUT:
        A structured dict (read-only receipt) with:
          * ``schema``            — DEFINITIVE_KNOWLEDGE_LEDGER_SCHEMA
          * ``priorityArms``      — list of {armIndex, knowledgeId, polarity,
                                     confidence, traumaLevel, coverage,
                                     priority} sorted DESC by priority
          * ``topKnowledgeIds``   — knowledge ids for the top_k arms (the
                                     curriculum surface the MitM scaffold reads)
        Knowledge ids are "" for unbound arms (legacy behavior).

    FAIL-OPEN: returns an empty receipt (schema + empty lists) on any error,
    or when the definitive-ledger gate is OFF (curriculum is inert until
    armed). Never raises.

    Args:
      state:             the trauma bank to read.
      coverage_payload:  optional {knowledge_id: coverage in [0, 1]}. Arms
                         whose knowledge id is absent default to coverage 0.
      top_k:             cap on the number of ranked arms returned.
    """

    receipt: dict[str, object] = {
        "schema": DEFINITIVE_KNOWLEDGE_LEDGER_SCHEMA,
        "priorityArms": [],
        "topKnowledgeIds": [],
    }
    try:
        if not _definitive_ledger_enabled_boundary():
            return receipt
        cov = coverage_payload if isinstance(coverage_payload, dict) else {}
        fails = state.fail_ema.clamp_min(0.0)
        peaks = state.peak_fail_ema.clamp_min(0.0)
        won = state.hard_won_ema.clamp_min(0.0)
        trauma_level = (fails + peaks + won)
        rows: list[dict[str, object]] = []
        for arm in range(state.num_arms):
            kid = ""
            try:
                if arm < len(state.arm_knowledge_ids):
                    kid = str(state.arm_knowledge_ids[arm] or "")
            except Exception:
                kid = ""
            pol = state.definitive_polarity.get(arm, DEFINITIVE_POLARITY_NONE)
            # Un-stamped arms default to confidence 1.0 so unverified hard
            # knowledge still ranks by trauma; stamped anchors use their
            # (>= floor) stamped confidence.
            conf = (
                float(state.definitive_confidence.get(arm, 1.0))
                if pol != DEFINITIVE_POLARITY_NONE
                else 1.0
            )
            coverage = 0.0
            if kid:
                raw_cov = cov.get(kid)
                if raw_cov is not None:
                    try:
                        coverage = max(0.0, min(1.0, float(raw_cov)))
                    except (TypeError, ValueError):
                        coverage = 0.0
            tl = float(trauma_level[arm].item())
            priority = tl * conf * max(
                0.0, 1.0 - coverage
            )
            rows.append(
                {
                    "armIndex": int(arm),
                    "knowledgeId": kid,
                    "polarity": pol,
                    "confidence": conf,
                    "traumaLevel": tl,
                    "coverage": coverage,
                    "priority": priority,
                }
            )
        def _priority_value(row: dict[str, object]) -> float:
            value = row["priority"]
            if not isinstance(value, (int, float)):
                raise TypeError("hard-knowledge priority is not numeric")
            return float(value)

        rows.sort(key=_priority_value, reverse=True)
        rows = rows[: max(0, int(top_k))]
        receipt["priorityArms"] = rows
        receipt["topKnowledgeIds"] = [
            str(r["knowledgeId"]) for r in rows if str(r["knowledgeId"])
        ]
        return receipt
    except Exception:  # pragma: no cover - fail-open
        return receipt


def definitive_knowledge_ledger(
    state: TensorTraumaState,
) -> dict[str, object]:
    """Read-only structured view: the definitive hard-won knowledge ledger.

    WHAT / WHY (doctrine):
        This IS "the definitive hard-won knowledge the model has / must learn".
        It maps each bound knowledge id -> {polarity, confidence, hard_won_ema,
        fail_ema, definitive, preserved}, so the operator can SEE which hard
        knowledge is positive-definitive (preserve), negative-definitive
        (avoid), and which is still a candidate (not yet verified-definitive).

    OUTPUT:
        A structured dict (read-only receipt) with:
          * ``schema``           — DEFINITIVE_KNOWLEDGE_LEDGER_SCHEMA
          * ``numArms``          — bank width
          * ``boundKnowledgeCount`` — number of arms with a knowledge binding
          * ``positiveDefinitiveCount`` / ``negativeDefinitiveCount`` /
            ``candidateCount`` — anchor-class breakdown
          * ``knowledge`` — list of {knowledgeId, armIndex, polarity,
                            confidence, hardWonEma, failEma, definitive,
                            preserved} (one row per bound arm, sorted by
                            confidence DESC then hard_won DESC)
        ``definitive`` is True only for arms stamped via ``mark_definitive``
        (verification-survived, high-confidence). ``preserved`` is True for
        positive-definitive arms with non-zero preservation weight (the
        anti-forgetting anchor set).

    FAIL-OPEN: returns an empty receipt (schema + zeros) on any error, or
    when the definitive-ledger gate is OFF. Never raises.
    """

    receipt: dict[str, object] = {
        "schema": DEFINITIVE_KNOWLEDGE_LEDGER_SCHEMA,
        "numArms": int(state.num_arms),
        "boundKnowledgeCount": 0,
        "positiveDefinitiveCount": 0,
        "negativeDefinitiveCount": 0,
        "candidateCount": 0,
        "knowledge": [],
    }
    try:
        preservation = positive_definitive_preservation_weights(state)
        rows: list[dict[str, object]] = []
        bound = 0
        pos_ct = 0
        neg_ct = 0
        cand_ct = 0
        for arm in range(state.num_arms):
            kid = ""
            try:
                if arm < len(state.arm_knowledge_ids):
                    kid = str(state.arm_knowledge_ids[arm] or "")
            except Exception:
                kid = ""
            if not kid:
                continue
            bound += 1
            pol = state.definitive_polarity.get(arm, DEFINITIVE_POLARITY_NONE)
            conf = float(state.definitive_confidence.get(arm, 0.0))
            definitive = pol != DEFINITIVE_POLARITY_NONE
            preserved = bool(
                definitive
                and pol == DEFINITIVE_POLARITY_POSITIVE
                and float(preservation[arm].item()) > 0.0
            )
            if definitive and pol == DEFINITIVE_POLARITY_POSITIVE:
                pos_ct += 1
            elif definitive and pol == DEFINITIVE_POLARITY_NEGATIVE:
                neg_ct += 1
            else:
                cand_ct += 1
            rows.append(
                {
                    "knowledgeId": kid,
                    "armIndex": int(arm),
                    "polarity": pol,
                    "confidence": conf,
                    "hardWonEma": float(state.hard_won_ema[arm].item()),
                    "failEma": float(state.fail_ema[arm].item()),
                    "definitive": bool(definitive),
                    "preserved": bool(preserved),
                }
            )
        def _ledger_order(row: dict[str, object]) -> tuple[float, float]:
            confidence = row["confidence"]
            hard_won = row["hardWonEma"]
            if not isinstance(confidence, (int, float)) or not isinstance(
                hard_won,
                (int, float),
            ):
                raise TypeError("definitive ledger order is not numeric")
            return float(confidence), float(hard_won)

        rows.sort(key=_ledger_order, reverse=True)
        receipt["boundKnowledgeCount"] = bound
        receipt["positiveDefinitiveCount"] = pos_ct
        receipt["negativeDefinitiveCount"] = neg_ct
        receipt["candidateCount"] = cand_ct
        receipt["knowledge"] = rows
        return receipt
    except Exception:  # pragma: no cover - fail-open
        return receipt


def trauma_state_snapshot(state: TensorTraumaState) -> dict[str, object]:
    """Full per-arm snapshot for persistence (one small JSON-safe dict).

    Used by the learn_loop boundary to persist trauma state next to the loop
    state. Returns plain Python lists so it is JSON-serializable. Round-trips
    through ``trauma_state_restore``. Kept tiny: 3 lists of ``num_arms``
    scalars plus config.

    ADDITIVE: the definitive-knowledge ledger fields (knowledge bindings,
    polarity stamps, confidence stamps) are persisted as small extra keys
    (``armKnowledgeIds``, ``definitivePolarity``, ``definitiveConfidence``)
    ONLY when non-empty, so an old snapshot round-trips unchanged and a new
    snapshot stays tiny when the ledger is unused. ``trauma_state_restore``
    treats all three keys as optional with legacy-empty defaults.
    """

    snap: dict[str, object] = {
        "schema": TRAUMA_SYSTEM_SCHEMA,
        "numArms": int(state.num_arms),
        "failDecay": float(state.fail_decay),
        "successDecay": float(state.success_decay),
        "failEma": state.fail_ema.detach().cpu().to(dtype=torch.float32).tolist(),
        "successEma": state.success_ema.detach().cpu().to(dtype=torch.float32).tolist(),
        "peakFailEma": state.peak_fail_ema.detach().cpu().to(dtype=torch.float32).tolist(),
        "hardWonEma": state.hard_won_ema.detach().cpu().to(dtype=torch.float32).tolist(),
        "knowledgeProEma": state.knowledge_pro_ema.detach().cpu().to(dtype=torch.float32).tolist(),
        "knowledgeAntiEma": state.knowledge_anti_ema.detach().cpu().to(dtype=torch.float32).tolist(),
        "behaviorProEma": state.behavior_pro_ema.detach().cpu().to(dtype=torch.float32).tolist(),
        "behaviorAntiEma": state.behavior_anti_ema.detach().cpu().to(dtype=torch.float32).tolist(),
        "verifiedExposureEma": state.verified_exposure_ema.detach().cpu().to(dtype=torch.float32).tolist(),
        "evidenceConfidenceEma": state.evidence_confidence_ema.detach().cpu().to(dtype=torch.float32).tolist(),
        "lastSuccessStep": state.last_success_step.detach().cpu().to(dtype=torch.long).tolist(),
        "globalStep": int(state.global_step.item()),
    }
    # Persist the additive ledger ONLY when populated (legacy snapshots stay
    # byte-identical; new snapshots only grow when the operator armed the
    # ledger). Knowledge ids are stored as a list (positional per arm).
    if state.arm_knowledge_ids:
        snap["armKnowledgeIds"] = list(str(k) for k in state.arm_knowledge_ids)
    if state.definitive_polarity:
        snap["definitivePolarity"] = {
            str(k): str(v) for k, v in state.definitive_polarity.items()
        }
    if state.definitive_confidence:
        snap["definitiveConfidence"] = {
            str(k): float(v) for k, v in state.definitive_confidence.items()
        }
    return snap


def trauma_state_restore(
    snapshot: dict[str, object] | None,
) -> TensorTraumaState | None:
    """Restore a ``TensorTraumaState`` from a snapshot, fail-open.

    Returns ``None`` if the snapshot is missing/malformed so the caller can
    fall back to a fresh state without raising. Adapted from the
    loss_telemetry fail-open boundary pattern.
    """

    if not isinstance(snapshot, dict):
        return None
    try:
        if snapshot.get("schema") != TRAUMA_SYSTEM_SCHEMA:
            return None
        num_arms_raw = snapshot.get("numArms")
        fail_decay_raw = snapshot.get(
            "failDecay",
            DEFAULT_TRAUMA_FAIL_DECAY,
        )
        success_decay_raw = snapshot.get(
            "successDecay",
            DEFAULT_TRAUMA_SUCCESS_DECAY,
        )
        if (
            not isinstance(num_arms_raw, int)
            or isinstance(num_arms_raw, bool)
            or not isinstance(fail_decay_raw, (int, float))
            or not isinstance(success_decay_raw, (int, float))
        ):
            return None
        num_arms = num_arms_raw
        fail_decay = float(fail_decay_raw)
        success_decay = float(success_decay_raw)
        state = TensorTraumaState(
            num_arms=num_arms,
            fail_decay=fail_decay,
            success_decay=success_decay,
        )
        fail_raw = snapshot.get("failEma")
        success_raw = snapshot.get("successEma")
        last_success_raw = snapshot.get("lastSuccessStep")
        if (
            not isinstance(fail_raw, list)
            or not isinstance(success_raw, list)
            or not isinstance(last_success_raw, list)
        ):
            return None
        state.fail_ema = torch.tensor(fail_raw, dtype=torch.float32)
        state.success_ema = torch.tensor(success_raw, dtype=torch.float32)
        peak_raw = snapshot.get("peakFailEma")
        if isinstance(peak_raw, list) and len(peak_raw) == num_arms:
            state.peak_fail_ema = torch.tensor(peak_raw, dtype=torch.float32)
        won_raw = snapshot.get("hardWonEma")
        if isinstance(won_raw, list) and len(won_raw) == num_arms:
            state.hard_won_ema = torch.tensor(won_raw, dtype=torch.float32)
        for snapshot_key, buffer_name in (
            ("knowledgeProEma", "knowledge_pro_ema"),
            ("knowledgeAntiEma", "knowledge_anti_ema"),
            ("behaviorProEma", "behavior_pro_ema"),
            ("behaviorAntiEma", "behavior_anti_ema"),
            ("verifiedExposureEma", "verified_exposure_ema"),
            ("evidenceConfidenceEma", "evidence_confidence_ema"),
        ):
            raw_buffer = snapshot.get(snapshot_key)
            if isinstance(raw_buffer, list) and len(raw_buffer) == num_arms:
                setattr(
                    state,
                    buffer_name,
                    torch.tensor(raw_buffer, dtype=torch.float32),
                )
        state.last_success_step = torch.tensor(
            last_success_raw,
            dtype=torch.long,
        )
        global_step_raw = snapshot.get("globalStep", 0)
        if not isinstance(global_step_raw, int) or isinstance(
            global_step_raw,
            bool,
        ):
            return None
        state.global_step = torch.tensor(global_step_raw, dtype=torch.long)
        # Width safety: if the persisted width disagrees with what the caller
        # expects, the caller is responsible for resizing; here we just
        # validate internal consistency.
        if (
            state.fail_ema.numel() != num_arms
            or state.success_ema.numel() != num_arms
            or state.peak_fail_ema.numel() != num_arms
            or state.hard_won_ema.numel() != num_arms
            or state.last_success_step.numel() != num_arms
        ):
            return None
        # ADDITIVE: restore the definitive-knowledge ledger when present.
        # All three keys are OPTIONAL with legacy-empty defaults, so an old
        # snapshot round-trips unchanged. Validated/clamped by __post_init__
        # via the constructor; here we pass them through so the restored
        # bank re-arms the ledger exactly as it was persisted.
        ids_raw = snapshot.get("armKnowledgeIds")
        if isinstance(ids_raw, list) and len(ids_raw) == num_arms:
            state.arm_knowledge_ids = tuple(str(k) for k in ids_raw)
            state.knowledge_to_arm = {
                str(kid): int(idx)
                for idx, kid in enumerate(state.arm_knowledge_ids)
                if str(kid).strip()
            }
        pol_raw = snapshot.get("definitivePolarity")
        if isinstance(pol_raw, dict):
            state.definitive_polarity = {
                int(k): str(v)
                for k, v in pol_raw.items()
                if str(v).strip()
                in {
                    DEFINITIVE_POLARITY_POSITIVE,
                    DEFINITIVE_POLARITY_NEGATIVE,
                    DEFINITIVE_POLARITY_NONE,
                }
                and 0 <= int(k) < num_arms
            }
        conf_raw = snapshot.get("definitiveConfidence")
        if isinstance(conf_raw, dict):
            state.definitive_confidence = {
                int(k): max(
                    DEFAULT_DEFINITIVE_CONFIDENCE_FLOOR,
                    min(DEFAULT_DEFINITIVE_CONFIDENCE_CEIL, float(v)),
                )
                for k, v in conf_raw.items()
                if 0 <= int(k) < num_arms
            }
        return state
    except Exception:
        return None


__all__ = [
    "DEFAULT_DEFINITIVE_CONFIDENCE",
    "DEFAULT_DEFINITIVE_CONFIDENCE_CEIL",
    "DEFAULT_DEFINITIVE_CONFIDENCE_FLOOR",
    "DEFAULT_HARD_KNOWLEDGE_FAIL_THRESHOLD",
    "DEFAULT_HARD_WON_EMA_THRESHOLD",
    "DEFAULT_HARD_WON_STRUGGLE_THRESHOLD",
    "DEFAULT_MUST_LEARN_COVERAGE_EPS",
    "DEFAULT_MUST_LEARN_PRIORITY_TOP_K",
    "DEFAULT_POSITIVE_DEFINITIVE_PRESERVATION_SCALE",
    "DEFAULT_TRAUMA_COOLDOWN_STEPS",
    "DEFAULT_TRAUMA_EPS",
    "DEFAULT_TRAUMA_FAIL_DECAY",
    "DEFAULT_TRAUMA_HIGH_FAIL_THRESHOLD",
    "DEFAULT_TRAUMA_MAX_FRACTION",
    "DEFAULT_TRAUMA_POSITIVE_SCALE",
    "DEFAULT_TRAUMA_PRESSURE_SCALE",
    "DEFAULT_TRAUMA_SUCCESS_DECAY",
    "DEFINITIVE_KNOWLEDGE_LEDGER_SCHEMA",
    "DEFINITIVE_POLARITY_CODE",
    "DEFINITIVE_POLARITY_DECODE",
    "DEFINITIVE_POLARITY_NEGATIVE",
    "DEFINITIVE_POLARITY_NONE",
    "DEFINITIVE_POLARITY_POSITIVE",
    "DEFINITIVE_UNBOUND_TOKEN",
    "ENV_TRAUMA_DEFINITIVE_LEDGER",
    "HARD_KNOWLEDGE_RECEIPT_SCHEMA",
    "TRAUMA_SYSTEM_SCHEMA",
    "TensorTraumaState",
    "apply_trauma_bias_to_quantile_router",
    "arm_for_knowledge",
    "bind_knowledge_to_arm",
    "definitive_hard_knowledge_surface_t",
    "definitive_knowledge_ledger",
    "hard_knowledge_must_learn_priority",
    "hard_knowledge_surface_receipt",
    "mark_definitive",
    "negative_hard_knowledge_arm_indices_t",
    "negative_hard_knowledge_mask_t",
    "positive_definitive_preservation_weights",
    "positive_hard_knowledge_arm_indices_t",
    "positive_hard_knowledge_mask_t",
    "positive_reinforcement_t",
    "trauma_bias_t",
    "trauma_pressure_t",
    "trauma_state_restore",
    "trauma_state_snapshot",
    "trauma_state_to_receipt",
    "update_trauma_from_outcome",
    "update_trauma_from_verified_outcome",
]