runtime/src/resynthesis/hard_knowledge_router_boundary.py

"""Hard-knowledge router boundary — bind trauma/MITM packet to live routers.

Hot path is **tensor-only**: ``HardKnowledgeSurfacePacket``, ``Tensor`` adjustments.
JSON/manifest adapters stay in ``coverage_pressure_boundary`` / assurance tools.

Schema: nnf.resynthesis.hard_knowledge_router_boundary.v1
"""

from __future__ import annotations

import os
import re
from dataclasses import dataclass
from typing import Any

import torch
import torch.nn as nn
from torch import Tensor

from resynthesis.hard_knowledge_surface import (
    HardKnowledgeSurfacePacket,
    hard_knowledge_surface_packet_t,
)
from resynthesis.trauma_system import TensorTraumaState

HARD_KNOWLEDGE_ROUTER_BOUNDARY_SCHEMA = (
    "nnf.resynthesis.hard_knowledge_router_boundary.v1"
)

DEFAULT_MITM_LOGIT_SCALE = 0.25
DEFAULT_HARD_KNOWLEDGE_PAGE_BOOST_FRACTION = 0.35


@dataclass(frozen=True, slots=True)
class StagedVerificationContext:
    """Typed observation boundary shared by training and staged verification."""

    available_layers: int
    trauma_hard_arms_t: Tensor | None


def _env_flag(name: str, default: bool = False) -> bool:
    raw = os.environ.get(name)
    if raw is None:
        return bool(default)
    return raw.strip().lower() in {"1", "true", "yes", "on"}


def hard_knowledge_router_enabled_boundary() -> bool:
    """Default ON when trauma system is enabled."""

    if _env_flag("NNF_HARD_KNOWLEDGE_ROUTER", default=False):
        return True
    if _env_flag("NNF_TRAUMA_SYSTEM", default=False):
        return True
    return _env_flag("NNF_RESYNTHESIS_TRAUMA_SYSTEM", default=False)


def branch_id_from_loop_id(loop_id: str) -> int:
    match = re.search(r"_b(\d+)", str(loop_id))
    if match:
        return int(match.group(1))
    match = re.search(r"gpu(\d+)", str(loop_id), flags=re.IGNORECASE)
    if match:
        return int(match.group(1))
    return 0


def load_mitm_coverage_payload_boundary(
    *,
    payload_path: str | os.PathLike[str] | None = None,
) -> object | None:
    """Load observer coverage payload for MITM hard-target ranking (boundary only)."""

    try:
        from pathlib import Path

        from resynthesis.coverage_pressure_boundary import (
            DEFAULT_LIVE_PAYLOAD_PATH,
            load_coverage_pressure_payload,
        )

        path = (
            Path(payload_path).expanduser()
            if payload_path is not None
            else DEFAULT_LIVE_PAYLOAD_PATH
        )
        payload = load_coverage_pressure_payload(path)
        if payload.undertrained_expert_ids or payload.undertrained_page_ids:
            return payload
    except Exception:
        pass
    return None


def page_catalog_ids_t_from_model(model: nn.Module) -> Tensor | None:
    """Best-effort read of the live NoNE page catalog tensor."""

    for module in model.modules():
        catalog = getattr(module, "page_catalog_ids_t", None)
        if isinstance(catalog, Tensor) and catalog.numel() > 0:
            return catalog.detach().reshape(-1).long()
    return None


def page_ids_to_catalog_arm_indices_t(
    model: nn.Module,
    page_ids: list[int] | Tensor,
) -> Tensor:
    """Map durable page IDs to catalog arm indices for the trauma bank."""

    if isinstance(page_ids, list):
        if not page_ids:
            return torch.empty(0, dtype=torch.long)
        ids_t = torch.tensor(page_ids, dtype=torch.long)
    else:
        ids_t = page_ids.detach().reshape(-1).long()
    if ids_t.numel() == 0:
        return ids_t
    catalog_t = page_catalog_ids_t_from_model(model)
    if catalog_t is None or catalog_t.numel() == 0:
        return ids_t[ids_t.ge(0)]
    identity_t = torch.arange(
        catalog_t.numel(),
        device=catalog_t.device,
        dtype=torch.long,
    )
    if bool(torch.equal(catalog_t, identity_t)):
        return ids_t[ids_t.ge(0) & ids_t.lt(catalog_t.numel())]
    expanded_ids_t = ids_t.reshape(-1, 1)
    matches_t = expanded_ids_t.eq(catalog_t.reshape(1, -1))
    has_match_t = matches_t.any(dim=1)
    if not bool(has_match_t.any().item()):
        return torch.empty(0, dtype=torch.long)
    positions_t = matches_t.to(dtype=torch.long).argmax(dim=1)
    return positions_t[has_match_t]


def lane_trauma_state_for_loop(
    loop_id: str,
    lane_bank: dict[str, Any] | None = None,
) -> TensorTraumaState | None:
    """Read the live per-lane trauma bank (learn_loop module scope)."""

    if lane_bank is None:
        try:
            from resynthesis import learn_loop as learn_loop_mod

            lane_bank = getattr(learn_loop_mod, "_LANE_TRAUMA_STATE", None)
        except Exception:
            return None
    if not isinstance(lane_bank, dict):
        return None
    state = lane_bank.get(str(loop_id))
    return state if isinstance(state, TensorTraumaState) else None


def _trauma_state_view_for_router(
    trauma_state: TensorTraumaState,
    num_arms: int,
) -> TensorTraumaState:
    """Non-mutating width-aligned view for a router (never truncates lane bank)."""

    if trauma_state.num_arms == num_arms:
        return trauma_state
    view = TensorTraumaState(
        num_arms=num_arms,
        fail_decay=trauma_state.fail_decay,
        success_decay=trauma_state.success_decay,
    ).to(device=trauma_state.fail_ema.device)
    prefix = min(trauma_state.num_arms, num_arms)
    with torch.no_grad():
        source_buffers = dict(trauma_state.named_buffers(recurse=False))
        for name, destination_t in view.named_buffers(recurse=False):
            source_t = source_buffers.get(name)
            if source_t is None:
                continue
            if destination_t.ndim == 0:
                destination_t.copy_(
                    source_t.to(
                        device=destination_t.device,
                        dtype=destination_t.dtype,
                    )
                )
            else:
                destination_t[:prefix].copy_(
                    source_t[:prefix].to(
                        device=destination_t.device,
                        dtype=destination_t.dtype,
                    )
                )
    view._knowledge_token_to_id = dict(trauma_state._knowledge_token_to_id)
    return view


def _align_trauma_width(
    trauma_state: TensorTraumaState,
    num_arms: int,
) -> TensorTraumaState:
    """Alias for router-local view (canonical lane bank is never mutated)."""

    return _trauma_state_view_for_router(trauma_state, num_arms)


def bind_hard_knowledge_to_quantile_router(
    router: nn.Module,
    *,
    trauma_state: TensorTraumaState | None,
    branch_id: int = 0,
    loop_id: str = "",
    coverage_payload: object | None = None,
) -> HardKnowledgeSurfacePacket | None:
    """Attach live hard-knowledge + MITM context to a ``QuantileBalancingRouter``."""

    if not hard_knowledge_router_enabled_boundary():
        return None
    num_arms = int(getattr(router, "num_experts", 0))
    if num_arms < 1 or trauma_state is None:
        return None
    try:
        aligned = _align_trauma_width(trauma_state, num_arms)
        packet = hard_knowledge_surface_packet_t(aligned, top_k=min(64, num_arms))
        # ``TensorTraumaState`` is an nn.Module, so assigning it here registers
        # the exact session/catalog bank beneath the router.  Its buffers now
        # participate in state_dict, device migration, and cold reload.
        setattr(router, "trauma_state", aligned)
        setattr(router, "_hard_knowledge_packet", packet)
        # Branch and loop identifiers are persistence-boundary metadata.  The
        # live router owns only the bound tensor state; strings/integers must
        # never be consulted to select a route or refresh its pressure.
        del branch_id, loop_id
        # Coverage payloads are observer evidence only.  They may be recorded
        # at the I/O boundary, but they never become route-logit authority.
        del coverage_payload
        return packet
    except Exception:
        return None


def refresh_hard_knowledge_packet(router: nn.Module) -> HardKnowledgeSurfacePacket | None:
    """Rebuild packet after trauma bank mutation (start of each route wave)."""

    trauma_state = getattr(router, "trauma_state", None)
    if not isinstance(trauma_state, TensorTraumaState):
        return getattr(router, "_hard_knowledge_packet", None)
    try:
        num_arms = int(getattr(router, "num_experts", trauma_state.num_arms))
        aligned = _align_trauma_width(trauma_state, num_arms)
        packet = hard_knowledge_surface_packet_t(aligned, top_k=min(64, num_arms))
        setattr(router, "trauma_state", aligned)
        setattr(router, "_hard_knowledge_packet", packet)
        setattr(router, "_mitm_hard_targets_t", None)
        return packet
    except Exception:
        return getattr(router, "_hard_knowledge_packet", None)


def model_owned_trauma_state_for_model(
    model: nn.Module,
) -> TensorTraumaState | None:
    """Return the first registered TRAUMA bank for boundary receipts only."""

    for module in model.modules():
        state = getattr(module, "trauma_state", None)
        if isinstance(state, TensorTraumaState):
            return state
    return None


def record_verified_hard_knowledge_outcome_to_model(
    model: nn.Module,
    *,
    knowledge_success_t: Tensor,
    behavior_success_t: Tensor,
    evidence_confidence_t: Tensor,
) -> Tensor:
    """Commit exact page-frontier evidence after verification.

    Parent paged routers retain the target-independent request that produced
    the completed transaction.  This function consumes its exact catalog
    positions and soft frontier weights, updates every selected page
    proportionally, and leaves the completed route unchanged.  The next route
    reads the newly committed model-owned buffers.
    """

    from resynthesis.trauma_system import (
        update_trauma_from_verified_outcome,
    )

    updated_count_t = knowledge_success_t.detach().new_zeros(
        (),
        dtype=torch.long,
    )
    seen_router_ids: set[int] = set()
    for module in model.modules():
        router = getattr(module, "quantile_router", None)
        request = getattr(module, "last_request", None)
        if not isinstance(router, nn.Module) or request is None:
            continue
        router_identity = id(router)
        if router_identity in seen_router_ids:
            continue
        state = getattr(router, "trauma_state", None)
        catalog_positions_t = getattr(
            request,
            "unique_page_catalog_positions_t",
            None,
        )
        frontier_weight_t = getattr(request, "frontier_weight_t", None)
        if (
            not isinstance(state, TensorTraumaState)
            or not isinstance(catalog_positions_t, Tensor)
            or not isinstance(frontier_weight_t, Tensor)
            or catalog_positions_t.ndim != 1
            or frontier_weight_t.ndim != 2
            or frontier_weight_t.shape[1] != catalog_positions_t.numel()
        ):
            continue
        seen_router_ids.add(router_identity)
        # Mean over rows preserves each request row's unit mass while avoiding
        # batch-size-dependent pressure.  Zero-mass/invalid columns are
        # rejected by the verified update boundary.
        page_weight_t = frontier_weight_t.detach().mean(dim=0)
        knowledge_t = knowledge_success_t.detach().reshape(()).to(
            device=state.fail_ema.device,
            dtype=state.fail_ema.dtype,
        )
        behavior_t = behavior_success_t.detach().reshape(()).to(
            device=state.fail_ema.device,
            dtype=state.fail_ema.dtype,
        )
        axis_count_t = state.fail_ema.new_tensor(2.0)
        pro_t = (knowledge_t + behavior_t) / axis_count_t
        anti_t = (
            (1.0 - knowledge_t) + (1.0 - behavior_t)
        ) / axis_count_t
        accepted_t = update_trauma_from_verified_outcome(
            state,
            arm_index_t=catalog_positions_t.detach().reshape(-1),
            frontier_weight_t=page_weight_t,
            pro_t=pro_t,
            anti_t=anti_t,
            evidence_confidence_t=evidence_confidence_t,
            knowledge_axis_t=state.fail_ema.new_ones(()),
            behavior_axis_t=state.fail_ema.new_ones(()),
        )
        updated_count_t = updated_count_t.to(device=accepted_t.device)
        updated_count_t = updated_count_t + accepted_t.numel()
        setattr(
            router,
            "_hard_knowledge_packet",
            hard_knowledge_surface_packet_t(
                state,
                top_k=state.num_arms,
                device=state.fail_ema.device,
            ),
        )
    return updated_count_t


def hard_knowledge_frontier_width_t(
    router: nn.Module,
    *,
    available: int,
    learned_k: int,
    uncertain: bool,
    width_t: Tensor,
) -> Tensor:
    """Widen frontier using hard-knowledge gap count (tensor-native)."""

    if not uncertain or not hard_knowledge_router_enabled_boundary():
        return width_t
    packet = getattr(router, "_hard_knowledge_packet", None)
    if not isinstance(packet, HardKnowledgeSurfacePacket):
        return width_t
    gap_count = packet.negative_count_t.to(dtype=width_t.dtype)
    if not bool(gap_count.gt(0).item()):
        return width_t
    from resynthesis.exploration_floor import exploration_route_width

    floor_width = exploration_route_width(
        available=available,
        learned_k=learned_k,
        uncertain=True,
    )
    boost = torch.clamp(
        gap_count,
        min=width_t.new_tensor(0.0),
        max=width_t.new_tensor(float(min(available, max(floor_width, 16)))),
    )
    return torch.maximum(width_t, boost.to(dtype=width_t.dtype))


def mitm_scaffold_logits_delta_t(
    router: nn.Module,
    *,
    num_candidates: int,
    device: torch.device,
    dtype: torch.dtype,
) -> Tensor:
    """Tensor-native temporary MILT bridge over the complete candidate axis.

    TRAUMA supplies real outcome difficulty and empirical learnability.
    Ownership progress attenuates the bridge to exactly zero after mastery.
    The router's learned activation fraction controls the bell range, so no
    host manifest, branch id, string gap key, fixed top-k, or answer target
    owns the route.
    """

    zero = torch.zeros(num_candidates, device=device, dtype=dtype)
    packet = getattr(router, "_hard_knowledge_packet", None)
    if not isinstance(packet, HardKnowledgeSurfacePacket):
        return zero
    if packet.num_arms != num_candidates:
        return zero
    difficulty_t = packet.negative_gap_signal_t.to(
        device=device,
        dtype=torch.float32,
    )
    if difficulty_t.numel() != num_candidates:
        return zero
    difficulty_unit_t = difficulty_t / difficulty_t.amax().clamp_min(1.0e-8)
    learnability_t = packet.empirical_learnability_t.to(
        device=device,
        dtype=torch.float32,
    ).clamp(0.0, 1.0)
    progress_t = packet.learning_progress_t.to(
        device=device,
        dtype=torch.float32,
    ).clamp(0.0, 1.0)
    ownership_gap_t = 1.0 - progress_t
    # A never-succeeded hard element still receives half-strength translation;
    # observed learnability raises that pressure, while proven ownership backs
    # it out exactly through ``ownership_gap_t``.
    bridge_strength_t = (
        difficulty_unit_t
        * ownership_gap_t
        * (0.5 + 0.5 * learnability_t)
    )
    activation_fraction = getattr(router, "activation_fraction", None)
    if not callable(activation_fraction):
        return zero
    learned_range_fraction_t = activation_fraction().reshape(()).to(
        device=device,
        dtype=torch.float32,
    ).clamp(0.0, 1.0)
    uncertainty_t = 1.0 - learnability_t
    range_t = (
        1.0
        + uncertainty_t
        * learned_range_fraction_t
        * max(1, num_candidates - 1)
    ).clamp_min(1.0e-3)
    positions_t = torch.arange(
        num_candidates,
        device=device,
        dtype=torch.float32,
    )
    distance_t = positions_t.unsqueeze(0) - positions_t.unsqueeze(1)
    bell_t = torch.exp(
        -0.5 * (distance_t / range_t.unsqueeze(1)).square()
    )
    # MILT is the translator between a failed hard endpoint and learnable
    # neighbouring paths; it must not hand the router the failed endpoint back
    # as its own answer.  Remove the diagonal while retaining the full learned
    # range around every hard element.
    # ``distance_t`` already carries the complete candidate-pair geometry.
    # Zero its exact diagonal without allocating a second dense floating-point
    # identity matrix.  Keep this out of place because ``bell_t`` participates
    # in the learned-range gradient.
    bell_t = bell_t.masked_fill(distance_t.eq(0), 0.0)
    exploration_t = (bridge_strength_t.unsqueeze(1) * bell_t).sum(dim=0)
    exploration_t = (
        exploration_t
        / exploration_t.amax().clamp_min(1.0e-8)
        * float(DEFAULT_MITM_LOGIT_SCALE)
    )
    active_t = bridge_strength_t.amax().gt(0.0).to(
        device=device,
        dtype=torch.float32,
    )
    return (exploration_t * active_t).to(dtype=dtype)


def hard_knowledge_page_logits_boost_t(
    *,
    packet: HardKnowledgeSurfacePacket,
    post_bias_logits_t: Tensor,
    score_span_t: Tensor,
) -> Tensor:
    """Apply signed pro-attraction/anti-suppression to page logits.

    The legacy function name is retained for callers, but negative definitive
    TRAUMA is never a positive boost.  The packet surface already encodes
    ``pro - anti`` in route-logit coordinates.
    """

    if post_bias_logits_t.numel() == 0:
        return post_bias_logits_t.new_zeros(post_bias_logits_t.shape)
    num_pages = post_bias_logits_t.shape[-1]
    if packet.definitive_surface_t.numel() != num_pages:
        raise ValueError("hard-knowledge packet differs from page catalog")
    active_signal_t = packet.definitive_surface_t.to(
        device=post_bias_logits_t.device,
        dtype=post_bias_logits_t.dtype,
    )
    peak = active_signal_t.abs().amax().clamp_min(
        active_signal_t.new_ones(()) * 1.0e-8
    )
    unit = active_signal_t / peak
    return (
        unit.unsqueeze(0)
        * score_span_t
        * (
            torch.ones_like(post_bias_logits_t).narrow(-1, 0, 1)
            * DEFAULT_HARD_KNOWLEDGE_PAGE_BOOST_FRACTION
        )
    )


def bind_hard_knowledge_to_model(
    model: nn.Module,
    *,
    loop_id: str,
    branch_id: int | None = None,
    trauma_state: TensorTraumaState | None = None,
    coverage_payload: object | None = None,
) -> int:
    """Bind hard-knowledge surface to every quantile router on ``model``."""

    if not hard_knowledge_router_enabled_boundary():
        return 0
    if trauma_state is None:
        trauma_state = lane_trauma_state_for_loop(loop_id)
    if branch_id is None:
        branch_id = branch_id_from_loop_id(loop_id)
    # Live coverage manifests remain observer-only.  They are intentionally
    # not loaded here and cannot alter model-owned route probabilities.
    bound = 0
    seen: set[int] = set()

    def _bind_router(router: nn.Module) -> None:
        nonlocal bound
        rid = id(router)
        if rid in seen:
            return
        seen.add(rid)
        active_state = trauma_state
        if active_state is None:
            resident_state = getattr(router, "trauma_state", None)
            active_state = (
                resident_state
                if isinstance(resident_state, TensorTraumaState)
                else None
            )
        if bind_hard_knowledge_to_quantile_router(
            router,
            trauma_state=active_state,
            branch_id=int(branch_id),
            loop_id=loop_id,
            coverage_payload=coverage_payload,
        ) is not None:
            bound += 1

    for module in model.modules():
        if module.__class__.__name__ == "QuantileBalancingRouter":
            _bind_router(module)
        router = getattr(module, "quantile_router", None)
        if isinstance(router, nn.Module):
            _bind_router(router)
    return bound


def staged_verification_context_for_model(
    model: nn.Module | None,
    *,
    loop_id: str = "",
) -> StagedVerificationContext:
    """Shared MITM/staged-verification context (layers + hard-knowledge arms)."""

    available_layers = 0
    trauma_hard_arms_t: Tensor | None = None
    if model is not None:
        try:
            science_stack = getattr(model, "science_stack", None)
            num_layers = getattr(science_stack, "num_layers", None)
            if isinstance(num_layers, int) and not isinstance(num_layers, bool):
                if num_layers > 0:
                    available_layers = num_layers
        except Exception:
            pass
    if loop_id:
        try:
            from resynthesis.mitm_trauma_learning_bridge import (
                hard_knowledge_arm_indices_t,
            )

            bank = lane_trauma_state_for_loop(loop_id)
            if bank is not None:
                catalog_t = (
                    page_catalog_ids_t_from_model(model) if model is not None else None
                )
                num_arms = (
                    int(catalog_t.numel())
                    if catalog_t is not None and catalog_t.numel() > 0
                    else bank.num_arms
                )
                view = _trauma_state_view_for_router(bank, num_arms)
                trauma_hard_arms_t = hard_knowledge_arm_indices_t(view)
        except Exception:
            trauma_hard_arms_t = None
    return StagedVerificationContext(
        available_layers=available_layers,
        trauma_hard_arms_t=trauma_hard_arms_t,
    )


__all__ = [
    "HARD_KNOWLEDGE_ROUTER_BOUNDARY_SCHEMA",
    "StagedVerificationContext",
    "bind_hard_knowledge_to_model",
    "bind_hard_knowledge_to_quantile_router",
    "branch_id_from_loop_id",
    "hard_knowledge_frontier_width_t",
    "hard_knowledge_page_logits_boost_t",
    "hard_knowledge_router_enabled_boundary",
    "lane_trauma_state_for_loop",
    "load_mitm_coverage_payload_boundary",
    "mitm_scaffold_logits_delta_t",
    "model_owned_trauma_state_for_model",
    "page_catalog_ids_t_from_model",
    "page_ids_to_catalog_arm_indices_t",
    "record_verified_hard_knowledge_outcome_to_model",
    "refresh_hard_knowledge_packet",
    "staged_verification_context_for_model",
]