runtime/src/resynthesis/causal_integration_tensor.py

"""Tensor-native integration of the capability stack WITH the Causal Algebra World Graph.

The model already runs ``CausalAlgebraWorldGraph`` (``resynthesis.causal_algebra``)
in its forward pass, producing a ``CausalTheoryProofPacket`` with per-action
disagreement, observation error, predicted outcomes, hypothesis posteriors, etc.
This module CONSUMES those tensor outputs and composes them with the tensor
capability modules (intent scorer, value function, calibration head) to produce
the four signals the model needs:

1. **Shaped action logits** — the model's raw logits + an exploration bonus
   derived from the causal packet's *disagreement* tensor (high disagreement =
   the model should explore that action) + a surprise signal from the
   *observation error*.
2. **Value estimate** — accumulated V(state) from the learned value table (TD
   updates come from the causal packet's observation outcomes).
3. **Intent composite** — the 8-axis route-quality shaping reward.
4. **Calibrated confidence** — the model's max-policy confidence rescaled by a
   learned temperature.

This is NOT a parallel stack — it is a thin composition layer that bridges the
causal spine's per-step tensor outputs to the long-term learning (value, credit)
and the exploration/reward shaping the model needs.  All learnable parameters
(value table, intent weights, bonus scale, calibration temperature) are
registered in the model's parameter set.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Final

import torch
import torch.nn as nn
import torch.nn.functional as F

from resynthesis.intent_scoring_tensor import TensorIntentScorer
from resynthesis.value_function_tensor import TensorValueFunction

CAUSAL_INTEGRATION_TENSOR_V1_SCHEMA = (
    "nnf.resynthesis.causal_integration_tensor.v1"
)
CAUSAL_INTEGRATION_TENSOR_SCHEMA = (
    "nnf.resynthesis.causal_integration_tensor.v5"
)
CAUSAL_INTEGRATION_TENSOR_CAUSAL_V1_LOCAL_STATE_NAMES: Final[
    frozenset[str]
] = (
    frozenset(
        {
            "exploration_bonus_scale",
            "intent_scorer.weights",
            "log_temperature",
            "surprise_weight",
            "value_function.value_table",
            "value_function.visit_counts",
        }
    )
)
CAUSAL_INTEGRATION_TENSOR_OUTCOME_V1_LOCAL_STATE_NAMES: Final[
    frozenset[str]
] = frozenset(
    {
        "exploration_meta.anti_fails",
        "exploration_meta.success_ema",
        "exploration_meta.visit_ema",
    }
)
CAUSAL_INTEGRATION_TENSOR_V1_LOCAL_STATE_NAMES: Final[frozenset[str]] = (
    CAUSAL_INTEGRATION_TENSOR_CAUSAL_V1_LOCAL_STATE_NAMES
    | CAUSAL_INTEGRATION_TENSOR_OUTCOME_V1_LOCAL_STATE_NAMES
)
CAUSAL_INTEGRATION_TENSOR_EXPLORATION_META_V2_LOCAL_STATE_NAMES: Final[
    frozenset[str]
] = frozenset(
    {
        "exploration_meta.confidence_threshold",
        "exploration_meta.exploit_mask",
        "exploration_meta.explore_mask",
        "exploration_meta.selection_logits",
    }
)
CAUSAL_INTEGRATION_TENSOR_ASSURANCE_V3_LOCAL_STATE_NAMES: Final[
    frozenset[str]
] = frozenset(
    {
        "ccl_loss.connectivity_logits",
        "ccl_loss.contrareactive_floor",
        "ccl_loss.contrareactive_weight",
        "consensus_clustering.correlation_threshold",
        "consensus_clustering.merge_temperature",
        "context_calibration.base_log_temperature",
        "context_calibration.context_to_temperature.bias",
        "context_calibration.context_to_temperature.weight",
        "knowledge_authority._collapse_floor_raw",
        "knowledge_authority._delta_floor_raw",
        "knowledge_authority._entropy_floor_raw",
        "knowledge_authority._nonzero_floor_raw",
        "knowledge_authority._norm_threshold_raw",
        "reasoning_authority._engagement_logit",
        "reasoning_authority._oscillation_logit",
        "reasoning_authority._stability_logit",
        "trauma_milt.history_len",
        "trauma_milt.milt.active",
        "trauma_milt.milt.milt_center",
        "trauma_milt.milt.milt_width",
        "trauma_milt.value_history",
    }
)
CAUSAL_INTEGRATION_TENSOR_KNOWLEDGE_TRANSFER_V4_LOCAL_STATE_NAMES: Final[
    frozenset[str]
] = frozenset(
    {
        # KnowledgeTransferTensor (verified fastest knowledge->weights): the
        # learnable cross-cluster transfer bank + anchor/distill temperatures
        # and sigmoid-gated weights. Added as a v4 capability-schema EXTENSION
        # so the versioned capability-seed migration registers + seeds it for
        # older checkpoints -- additive: new elements extend the schema, they
        # never block a checkpoint load (the corruption guards stay intact).
        "knowledge_transfer.transfer_bank.cluster_affinity",
        "knowledge_transfer.transfer_bank.capability_routing",
        "knowledge_transfer.transfer_bank.verifier_affinity",
        "knowledge_transfer.transfer_bank.transfer_proj_down.weight",
        "knowledge_transfer.transfer_bank.transfer_proj_up.weight",
        "knowledge_transfer.transfer_bank.transfer_scale",
        "knowledge_transfer.anchor_log_temperature",
        "knowledge_transfer.distill_log_temperature",
        "knowledge_transfer.anchor_weight_raw",
        "knowledge_transfer.distill_weight_raw",
        "knowledge_transfer.contact_seek.contact_log_temperature",
        "knowledge_transfer.contact_seek.contact_weight_raw",
        "knowledge_transfer.contact_seek.adaptive.top_n_log",
        "knowledge_transfer.contact_seek.adaptive.band_width_log",
        "knowledge_transfer.contact_seek.adaptive.correlation_logit",
        "knowledge_transfer.contact_seek.adaptive.region_scale_log",
        "knowledge_transfer.contact_seek.adaptive.min_contact_logit",
        "knowledge_transfer.contact_seek.adaptive.min_diff_logit",
        "knowledge_transfer.contact_seek.adaptive.region_reposition_raw",
    }
)
CAUSAL_INTEGRATION_TENSOR_MHC_V5_LOCAL_STATE_NAMES: Final[
    frozenset[str]
] = frozenset(
    {
        "ccl_loss.mhc_head.bias",
        "ccl_loss.mhc_head.ds_weight",
        "ccl_loss.mhc_head.mix",
        "ccl_loss.mhc_head.weight",
        "intent_scorer.mhc_residual_head.alpha_head.bias",
        "intent_scorer.mhc_residual_head.alpha_head.ds_weight",
        "intent_scorer.mhc_residual_head.alpha_head.mix",
        "intent_scorer.mhc_residual_head.alpha_head.weight",
        "intent_scorer.mhc_residual_head.delta_head.bias",
        "intent_scorer.mhc_residual_head.delta_head.ds_weight",
        "intent_scorer.mhc_residual_head.delta_head.mix",
        "intent_scorer.mhc_residual_head.delta_head.weight",
        "intent_scorer.mhc_residual_head.norm.weight",
    }
)
CAUSAL_INTEGRATION_TENSOR_LOCAL_STATE_NAMES: Final[frozenset[str]] = (
    CAUSAL_INTEGRATION_TENSOR_V1_LOCAL_STATE_NAMES
    | CAUSAL_INTEGRATION_TENSOR_EXPLORATION_META_V2_LOCAL_STATE_NAMES
    | CAUSAL_INTEGRATION_TENSOR_ASSURANCE_V3_LOCAL_STATE_NAMES
    | CAUSAL_INTEGRATION_TENSOR_KNOWLEDGE_TRANSFER_V4_LOCAL_STATE_NAMES
    | CAUSAL_INTEGRATION_TENSOR_MHC_V5_LOCAL_STATE_NAMES
)


@dataclass
class CausalIntegrationOutput:
    """The signals the model consumes from the capability stack."""

    shaped_action_logits: torch.Tensor  # [batch, num_actions]
    value: torch.Tensor  # [batch]
    intent_composite: torch.Tensor | None  # [batch] or None
    calibrated_confidence: torch.Tensor  # [batch]
    exploration_bonus: torch.Tensor  # [batch, num_actions]
    surprise: torch.Tensor  # [batch] — mean observation error from the causal packet
    # Enhanced signals from the doctrine-hardened modules (all optional / None when
    # the corresponding input was not provided — backwards-compatible).
    auxiliary_loss: torch.Tensor | None = None  # scalar — CCL connectivity-contrastive loss
    causal_auxiliary_loss: torch.Tensor | None = None
    transferred_hidden: torch.Tensor | None = None
    reasoning_diagnostic: object | None = None  # ReasoningDiagnostic from the RLA
    knowledge_verified: object | None = None  # KnowledgeVerification from the KLA


class CausalIntegrationTensor(nn.Module):
    """Composition layer bridging the Causal Algebra World Graph to capability modules.

    Instantiate as a model submodule.  Call ``forward`` each step with the
    causal packet's tensor outputs + the model's action logits.  The causal
    spine is NOT duplicated here — the model already runs it; this module
    consumes its proof packet.
    """

    def __init__(
        self,
        *,
        hidden_size: int,
        transfer_dim: int,
        hypothesis_count: int = 8,
        table_size: int = 4096,
        device: torch.device | None = None,
        dtype: torch.dtype = torch.float32,
    ) -> None:
        super().__init__()
        self.table_size = table_size
        self._dtype = dtype

        # Learnable submodules (value table + intent scorer).
        self.value_function = TensorValueFunction(
            table_size=table_size, device=device, dtype=dtype,
        )
        self.intent_scorer = TensorIntentScorer(
            mhc_residual=True,
            device=device,
            dtype=dtype,
        )

        # Exploration bonus scale (learnable — how much the causal disagreement
        # shapes the model's action logits).
        self.exploration_bonus_scale = nn.Parameter(torch.tensor(0.1, dtype=dtype))

        # Surprise weighting (how much observation error suppresses confidence).
        self.surprise_weight = nn.Parameter(torch.tensor(0.5, dtype=dtype))

        # Calibration temperature (self-models confidence reliability).
        self.log_temperature = nn.Parameter(torch.tensor(0.0, dtype=dtype))

        from resynthesis.exploration_meta_controller_tensor import TensorMetaController

        self.exploration_meta = TensorMetaController(device=device, dtype=dtype)

        # --- Doctrine-hardened + ContactSeek-enhanced submodules (all tensor-native) ---
        from resynthesis.ccl_loss_tensor import CCLLoss
        from resynthesis.knowledge_layer_authority_tensor import KnowledgeLayerAuthority
        from resynthesis.reasoning_layer_authority_tensor import ReasoningLayerAuthority
        from resynthesis.differential_consensus_tensor import DifferentialConsensusClustering
        from resynthesis.context_calibration_tensor import ContextCalibrationHead
        from resynthesis.trauma_milt_scaffolding_tensor import (
            HillClimbScaffoldingCycle,
        )

        # CCL connectivity-contrastive loss for causal DAG recovery.
        self.ccl_loss = CCLLoss(
            num_hypotheses=hypothesis_count,
            mhc_projection=True,
            device=device,
            dtype=dtype,
        )
        # KLA — knowledge verification (called externally during grading).
        self.knowledge_authority = KnowledgeLayerAuthority()
        # RLA — real-time reasoning diagnostic.
        self.reasoning_authority = ReasoningLayerAuthority(device=device, dtype=dtype)
        # Differential consensus clustering for disagreement → causal regions.
        self.consensus_clustering = DifferentialConsensusClustering()
        # Context-conditional calibration (replaces simple temperature when context available).
        self.context_calibration = ContextCalibrationHead(
            context_dim=8,  # matches intent axes
            device=device,
            dtype=dtype,
        )
        # The 16-item axis is the deterministic v26 checkpoint geometry, not
        # a routing/top-k ceiling. Runtime exploration remains model-owned and
        # later accepted generations may grow through a versioned migration.
        self.trauma_milt = HillClimbScaffoldingCycle(
            num_items=16,
            device=device,
            dtype=dtype,
        )
        # Knowledge-transfer coordinator (verified fastest knowledge->weights):
        # forward-KL anchor + soft-target distillation + transfer residual
        # (Approach B: hosts KnowledgeTransferBank in the gradient path).
        # Registered as the v4 capability-schema extension; the versioned
        # migration seeds it for older checkpoints.  Both dimensions are
        # required from the accepted expandable graph: defaults such as
        # 4096/256 would silently turn today's parent seam and transfer rank
        # into tomorrow's capacity ceiling.
        from resynthesis.knowledge_transfer_tensor import KnowledgeTransferTensor

        self.knowledge_transfer = KnowledgeTransferTensor(
            hidden_size=hidden_size,
            transfer_dim=transfer_dim,
            device=device,
            dtype=dtype,
        )

    def bind_anti_thompson_registry(self, registry: object) -> None:
        """Optional anti-Thompson fail bank for quantile routers on science stack."""

        self._anti_thompson_registry = registry

    def rebuild_nonpersistent_buffers(self) -> None:
        """Materialize deterministic runtime state after meta-device cold load."""

        from resynthesis.ccl_loss_tensor import dag_acyclicity_mask

        connectivity = self.ccl_loss.connectivity_logits
        self.ccl_loss.acyclicity_mask = dag_acyclicity_mask(
            self.ccl_loss.num_hypotheses,
            upper_triangular=self.ccl_loss.upper_triangular,
            device=connectivity.device,
            dtype=connectivity.dtype,
        )
        self.knowledge_transfer.rebuild_nonpersistent_buffers()

    # -- forward (the model calls this each step with the causal packet) --

    def forward(
        self,
        *,
        action_logits: torch.Tensor,
        causal_disagreement: torch.Tensor | None = None,
        observation_error: torch.Tensor | None = None,
        predicted_outcomes: torch.Tensor | None = None,
        posterior: torch.Tensor | None = None,
        state_keys: list[str] | None = None,
        intent_axes: torch.Tensor | None = None,
        student_logits: torch.Tensor | None = None,
        prior_additive_logits: torch.Tensor | None = None,
        learned_teacher_logits: torch.Tensor | None = None,
        student_hidden: torch.Tensor | None = None,
        prior_additive_hidden: torch.Tensor | None = None,
        prompt_len: int = 0,
        contact_feature_stack: torch.Tensor | None = None,
        page_count: int = 0,
        source_cluster_index_t: torch.Tensor | None = None,
        target_cluster_index_t: torch.Tensor | None = None,
        transfer_step_t: torch.Tensor | None = None,
    ) -> CausalIntegrationOutput:
        """Compose the causal spine's outputs into the four model signals.

        Parameters
        ----------
        action_logits : ``[batch, num_actions]``
            The model's raw action logits (from its policy head).
        causal_disagreement : ``[batch, num_actions] | None``
            Per-action disagreement from the causal packet's falsification step
            (``falsifying_experiment.disagreement_t``).  High = the hypotheses
            disagree on that action → the model should explore it.
        observation_error : ``[batch, num_hypotheses] | None``
            Per-hypothesis observation error from the causal packet
            (``observation_error_t``).  Reduced to a per-batch surprise scalar.
        state_keys : ``list[str] | None``
            State identifiers for value-table lookup (hashed to indices).
        intent_axes : ``[batch, 8] | None``
            The 8-axis route-quality measurement (optional shaping reward).
        """

        batch_size, num_actions = action_logits.shape

        # --- Exploration bonus from the causal packet's disagreement ---
        if causal_disagreement is not None:
            # Normalize disagreement to [0, 1] per batch row (relative scale).
            max_disagreement = causal_disagreement.amax(dim=-1, keepdim=True).clamp_min(1e-8)
            normalized_disagreement = causal_disagreement / max_disagreement
            consensus = self.consensus_clustering(causal_disagreement)
            region_mass = consensus.soft_memberships.sum(dim=1).clamp_min(1e-8)
            region_mean = consensus.region_deltas / region_mass
            consensus_action = torch.bmm(
                consensus.soft_memberships,
                region_mean.unsqueeze(-1),
            ).squeeze(-1)
            consensus_scale = consensus_action.abs().amax(
                dim=-1,
                keepdim=True,
            ).clamp_min(1e-8)
            normalized_consensus = consensus_action / consensus_scale
            normalized_disagreement = 0.5 * (
                normalized_disagreement + normalized_consensus
            )
            bonus = normalized_disagreement * torch.tanh(self.exploration_bonus_scale)
            visit_counts = self.value_function.visit_counts
            if visit_counts.numel() >= batch_size:
                from resynthesis.anti_systems_bridge import widen_exploration_bonus_with_meta

                values = self.value_function.evaluate_many(
                    state_keys if state_keys is not None else [""] * batch_size
                )
                bonus = widen_exploration_bonus_with_meta(
                    base_bonus_t=bonus,
                    visit_counts_t=visit_counts[:batch_size],
                    values_t=values,
                    max_abs_value=values.abs().amax().clamp_min(1.0e-8),
                    meta_controller=self.exploration_meta,
                )
        else:
            bonus = torch.zeros(batch_size, num_actions, device=action_logits.device, dtype=action_logits.dtype)

        # --- Surprise from the causal packet's observation error ---
        if observation_error is not None:
            surprise = observation_error.mean(dim=-1)  # [B]
        else:
            surprise = torch.zeros(batch_size, device=action_logits.device, dtype=action_logits.dtype)

        # Shaped logits: raw logits + exploration bonus - surprise dampening.
        shaped_logits = action_logits + bonus - self.surprise_weight * surprise.unsqueeze(-1) * 0.1

        # --- Value estimate (accumulated V from the value table) ---
        if state_keys is not None:
            value = self.value_function.evaluate_many(state_keys)  # [B]
        else:
            value = torch.zeros(batch_size, device=action_logits.device, dtype=action_logits.dtype)

        # --- Intent composite (shaping reward, if observations provided) ---
        intent = None
        if intent_axes is not None:
            intent = self.intent_scorer.composite(intent_axes)  # [B]

        # --- Calibrated confidence ---
        confidence = torch.softmax(shaped_logits, dim=-1).max(dim=-1).values
        calibrated = self.calibration_scale(confidence)
        if causal_disagreement is not None and observation_error is not None:
            context_source = torch.cat(
                (causal_disagreement, observation_error),
                dim=-1,
            )
            context_features = F.adaptive_avg_pool1d(
                context_source.unsqueeze(1),
                self.context_calibration.context_dim,
            ).squeeze(1)
            calibrated = self.context_calibration(
                calibrated,
                context_features,
            )

        # --- CCL connectivity-contrastive loss (causal DAG recovery) ---
        # CCL must learn from the executable causal graph's own intervention
        # outcomes and posterior. Reconstructing a broadcast tensor from
        # disagreement/error destroyed hypothesis/action/world structure and
        # trained a synthetic relation that the graph never predicted.
        if (predicted_outcomes is None) != (posterior is None):
            raise ValueError(
                "causal predicted outcomes and posterior must be supplied together"
            )
        causal_aux_loss = None
        if predicted_outcomes is not None and posterior is not None:
            ccl = self.ccl_loss(predicted_outcomes, posterior)
            anti = self.ccl_loss.contrareactive_loss(predicted_outcomes)
            causal_aux_loss = ccl + anti
        aux_loss = causal_aux_loss

        # --- RLA real-time reasoning diagnostic ---
        reasoning_diag = self.reasoning_authority(
            posterior=posterior,
            disagreement=causal_disagreement,
            observation_error=observation_error,
            expert_activations=bonus,
            stop_scores=(
                shaped_logits[:, :3] if shaped_logits.shape[-1] >= 3 else None
            ),
        )

        # Additive knowledge transfer is fail-closed on authority or geometry
        # drift.  The previous broad exception handler silently discarded a
        # malformed reference and could hide an accidental parent-logit anchor.
        # Generic reference/teacher names are deliberately absent: callers may
        # supply only prior accepted additive state or a genuinely learned
        # teacher.
        kt_student = student_logits if student_logits is not None else action_logits
        kt_out = self.knowledge_transfer(
            student_logits=kt_student,
            prior_additive_logits=prior_additive_logits,
            learned_teacher_logits=learned_teacher_logits,
            hidden=student_hidden,
            prior_additive_hidden=prior_additive_hidden,
            prompt_len=prompt_len,
            contact_feature_stack=contact_feature_stack,
            page_count=page_count,
            source_cluster_index_t=source_cluster_index_t,
            target_cluster_index_t=target_cluster_index_t,
            step_t=transfer_step_t,
        )
        if kt_out.auxiliary_loss is not None:
            aux_loss = (
                kt_out.auxiliary_loss
                if aux_loss is None
                else aux_loss + kt_out.auxiliary_loss
            )

        return CausalIntegrationOutput(
            shaped_action_logits=shaped_logits,
            value=value,
            intent_composite=intent,
            calibrated_confidence=calibrated,
            exploration_bonus=bonus,
            surprise=surprise,
            # Enhanced signals:
            auxiliary_loss=aux_loss,
            causal_auxiliary_loss=causal_aux_loss,
            transferred_hidden=kt_out.transferred_hidden,
            reasoning_diagnostic=reasoning_diag,
            knowledge_verified=None,  # KLA called externally during grading
        )

    # -- online update (called after each transition) --------------------

    def update_from_transition(
        self,
        *,
        state_keys: list[str],
        next_state_keys: list[str],
        rewards: torch.Tensor,
        alpha: float = 0.1,
        discount: float = 0.9,
    ) -> torch.Tensor:
        """TD(0) value update from a transition. Returns per-batch TD errors."""

        return self.value_function.apply_td_update_batch(
            state_keys=state_keys,
            next_state_keys=next_state_keys,
            rewards=rewards,
            alpha=alpha,
            discount=discount,
        )

    # -- confidence calibration (learned temperature scaling) -----------

    def calibration_scale(self, confidence: torch.Tensor) -> torch.Tensor:
        """Scale confidence by a learned temperature (self-modeling)."""

        temperature = torch.exp(self.log_temperature)
        return confidence.clamp(1e-8, 1.0) ** (1.0 / temperature.clamp_min(1e-8))

    # -- goal progress (value as fraction toward the merit endstate) ----

    def goal_progress(self, state_keys: list[str]) -> torch.Tensor:
        """Value of each state as progress [0,1] toward the merit endstate."""

        values = self.value_function.evaluate_many(state_keys)
        return values.clamp(0.0, 1.0)

    def parameter_count(self) -> int:
        return sum(p.numel() for p in self.parameters() if p.requires_grad)

    # -- KLA knowledge verification (called externally during grading) ----

    def verify_page_knowledge(
        self,
        page_weights: torch.Tensor,
        *,
        page_id: str = "unknown",
    ) -> object:
        """Verify a page has real trained knowledge (KLA)."""

        return self.knowledge_authority.verify_page_knowledge(
            page_weights, page_id=page_id,
        )

    def verify_model_knowledge(
        self,
        model_state_dict: dict[str, torch.Tensor],
        *,
        sampled_pages: int = 16,
    ) -> object:
        """Sample pages from the model and verify knowledge retention (KLA)."""

        return self.knowledge_authority.verify_model_knowledge(
            model_state_dict, sampled_pages=sampled_pages,
        )