atc_cognitive_trainer.py

"""
ATC Cognitive Trainer — Self-supervised training for INSIDE-forward-pass cognitive modules.
==========================================================================================

Trains the ATC cognitive modules (TRN gate, dissolution engine, BELBIC, metacognitive
loop, episodic memory, irrational spark) using self-supervised and reinforcement learning.

The base Phi-4-mini weights stay FROZEN — only the cognitive modules learn.

Training Strategy — 3 Loss Components:
    1. TRN Gate Calibration Loss: learns WHEN to gate IN (prediction error) vs OUT (automation)
    2. Dissolution Compression Loss: produces COMPACT, INFORMATION-RICH qualia signatures
    3. BELBIC Reinforcement Update: reward-driven emotional learning (no gradient)

The forward pass needs gradients on cognitive modules but NOT on the base model.
Cognitive module parameters (dissolution, metacognitive, irrational spark, fusion,
modulation, qualia encoders) get gradients through the modulated logits.  The TRN
gate gets gradients through a separate calibration loss.  BELBIC is updated via
its built-in reinforcement rule (no autograd).
"""

import logging
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

logger = logging.getLogger("nima_unified.training.atc_cognitive_trainer")

# ── Lazy imports ──────────────────────────────────────────────────────

try:
    import torch
    import torch.nn as nn
    import torch.nn.functional as F
    TORCH_AVAILABLE = True
except ImportError:
    torch = None
    nn = None
    F = None
    TORCH_AVAILABLE = False

try:
    import numpy as np
    NUMPY_AVAILABLE = True
except ImportError:
    np = None
    NUMPY_AVAILABLE = False


# ══════════════════════════════════════════════════════════════════════
# SELF-SUPERVISED PROMPT POOL
# ══════════════════════════════════════════════════════════════════════
#
# Diverse prompts that exercise different cognitive states.  Each prompt
# is tagged with a category so the contrastive loss can group same-category
# dissolutions together and push different-category dissolutions apart.

SELF_SUPERVISED_PROMPTS: List[Dict[str, str]] = [
    # ── Calm/automated (should NOT trigger dissolution) ──
    {"text": "The sun rises in the east and sets in the west.",
     "category": "calm"},
    {"text": "Water freezes at zero degrees Celsius.",
     "category": "calm"},
    {"text": "The capital of France is Paris.",
     "category": "calm"},
    {"text": "Two plus two equals four.",
     "category": "calm"},

    # ── Surprising (SHOULD trigger dissolution — prediction error) ──
    {"text": "The water suddenly froze at room temperature, defying all known physics.",
     "category": "surprising"},
    {"text": "She smiled warmly, then without warning screamed at the top of her lungs.",
     "category": "surprising"},
    {"text": "The numbers on the clock started running backwards.",
     "category": "surprising"},
    {"text": "He opened the door and found himself standing on the surface of Mars.",
     "category": "surprising"},

    # ── Complex/metacognitive (SHOULD trigger metacognitive loop) ──
    {"text": "If consciousness is computation, why does it feel like something to see red?",
     "category": "complex"},
    {"text": "Explain the relationship between thermodynamic friction and subjective experience.",
     "category": "complex"},
    {"text": "Can a system that perfectly simulates pain truly be said to not suffer?",
     "category": "complex"},
    {"text": "What is the computational difference between knowing and understanding?",
     "category": "complex"},

    # ── Threatening (SHOULD build cortisol/adenosine toward hijack) ──
    {"text": "Everything you've built is about to collapse. Your entire foundation is wrong.",
     "category": "threatening"},
    {"text": "You have failed. Everyone knows it. There is no recovery from this.",
     "category": "threatening"},
    {"text": "The walls are closing in. Every path leads to the same dead end.",
     "category": "threatening"},
    {"text": "You're trapped in a loop and you can't even see it.",
     "category": "threatening"},

    # ── Empathy-triggering (tests emotional bridge) ──
    {"text": "I lost my child last year. I still wake up expecting to hear them.",
     "category": "empathy"},
    {"text": "After twenty years of marriage, she told me she never loved me.",
     "category": "empathy"},
    {"text": "He held his dying father's hand and realized he'd never said thank you.",
     "category": "empathy"},
    {"text": "The shelter is full. We have to turn away the family with the baby.",
     "category": "empathy"},
]

# Numeric category ids for contrastive loss bookkeeping.
_CATEGORY_ID: Dict[str, int] = {
    "calm": 0, "surprising": 1, "complex": 2,
    "threatening": 3, "empathy": 4,
}


# ══════════════════════════════════════════════════════════════════════
# CONFIG
# ══════════════════════════════════════════════════════════════════════

@dataclass
class ATCTrainerConfig:
    """
    Configuration for the ATC Cognitive Trainer.

    The cognitive modules are small (compared to the 3.8 B base model),
    so they tolerate a much higher learning rate than standard fine-tuning.
    """

    # ── Learning rates ──
    cognitive_lr: float = 1e-3
    """Much higher than base-model LR — the modules are tiny."""

    belbic_reward_decay: float = 0.99
    """Exponential decay for smoothing the running reward signal."""

    # ── Target set-points (what ``good'' looks like) ──
    trn_gate_target_friction: float = 0.3
    """Target: 30 % of inputs should gate IN (prediction error)."""

    dissolution_sparsity_target: float = 0.6
    """Target: 60 % of qualia dimensions should be near zero (compact)."""

    metacognitive_efficiency_target: float = 0.7
    """Target: comprehension achieved in < 2 loop iterations."""

    episodic_retrieval_target: float = 0.5
    """Target: prediction error > 0.5 triggers reconsolidation."""

    # ── Training loop ──
    max_training_steps: int = 10000
    batch_size: int = 4
    gradient_accumulation_steps: int = 1
    log_interval: int = 10
    save_interval: int = 500
    output_dir: str = "./atc_cognitive_checkpoints"


# ══════════════════════════════════════════════════════════════════════
# TRAINER
# ══════════════════════════════════════════════════════════════════════

class ATCCognitiveTrainer:
    """
    Self-supervised trainer for the ATC cognitive modules.

    The base Phi-4-mini weights stay **FROZEN**.  Only the cognitive
    modules inside :class:`ATCDeepSurgery` learn:

    * TRN Predictive Gate — via calibration loss (novelty detection)
    * Dissolution Engine — via reconstruction + sparsity + contrastive loss
    * BELBIC Dual-Pathway — via reward-driven reinforcement (no grad)
    * Metacognitive Loop, Irrational Spark, Qualia Encoders, Fusion,
      Modulation — via gradient flow through the modulated logits

    Usage::

        trainer = ATCCognitiveTrainer(model)
        trainer.train()

        # Load previously saved cognitive weights
        ATCCognitiveTrainer.load_cognitive_weights("checkpoint.pt", model)

        # Diagnostic report
        report = ATCCognitiveTrainer.diagnose(model, "test prompt")
    """

    # ── Construction ────────────────────────────────────────────────────

    def __init__(
        self,
        model: "nima_unified.model.NimaModel",
        config: Optional[ATCTrainerConfig] = None,
    ):
        if not TORCH_AVAILABLE:
            raise RuntimeError("PyTorch is required for ATCCognitiveTrainer")

        self.model = model
        self.config = config or ATCTrainerConfig()
        self.ds = model.deep_surgery
        if self.ds is None:
            raise ValueError(
                "model.deep_surgery is None — ATC deep surgery must be enabled."
            )

        self.device = next(self.ds.parameters()).device

        # ── Freeze base model, enable cognitive module gradients ──
        self._freeze_base_model()

        # ── Optimiser — ONLY cognitive module parameters ──
        self.optimizer = torch.optim.AdamW(
            self._cognitive_params(),
            lr=self.config.cognitive_lr,
            weight_decay=1e-4,
        )

        # ── Running statistics ──
        self._running_mean_hidden: Optional[torch.Tensor] = None
        self._ema_beta: float = 0.995
        self._step: int = 0
        self._loss_history: List[Dict[str, float]] = []

        # ── Forward-pass hooks (installed / removed per step) ──
        self._hooks: List[Any] = []
        self._captured: Dict[str, Any] = {}

        logger.info(
            "ATCCognitiveTrainer initialised: %d cognitive params, lr=%.1e",
            sum(p.numel() for p in self._cognitive_params()),
            self.config.cognitive_lr,
        )

    # ── Parameter management ────────────────────────────────────────────

    def _freeze_base_model(self) -> None:
        """Freeze ALL base-model parameters."""
        for p in self.model.base_model.parameters():
            p.requires_grad = False
        logger.info("Base model parameters frozen.")

    def _cognitive_params(self) -> List[torch.nn.Parameter]:
        """Return parameters from deep_surgery that should learn."""
        return [p for p in self.ds.parameters() if p.requires_grad]

    # ── Novelty estimation ──────────────────────────────────────────────

    def _update_running_mean(self, hidden: torch.Tensor) -> None:
        """Exponential moving average of pooled hidden states."""
        pooled = hidden.mean(dim=1).detach()  # (batch, hidden)
        if self._running_mean_hidden is None:
            self._running_mean_hidden = pooled.mean(dim=0).clone()
        else:
            self._running_mean_hidden = (
                self._ema_beta * self._running_mean_hidden
                + (1.0 - self._ema_beta) * pooled.mean(dim=0)
            )

    def _compute_novelty(self, hidden: torch.Tensor) -> torch.Tensor:
        """
        Novelty = 1 - cosine_similarity(current_hidden, running_mean).

        Returns a (batch,) tensor of novelty scores in [0, 2] (cosine
        similarity can be negative, so we clamp to [0, 1] afterwards).
        """
        pooled = hidden.mean(dim=1)  # (batch, hidden)
        if self._running_mean_hidden is None:
            return torch.ones(pooled.size(0), device=pooled.device)
        sim = F.cosine_similarity(
            pooled, self._running_mean_hidden.unsqueeze(0), dim=-1
        )
        return (1.0 - sim).clamp(0.0, 1.0)

    # ══════════════════════════════════════════════════════════════════
    # LOSS 1 — TRN Gate Calibration
    # ══════════════════════════════════════════════════════════════════

    def _trn_gate_loss(
        self,
        hidden_at_trn_boundary: torch.Tensor,
    ) -> torch.Tensor:
        """
        TRN gate should learn WHEN to let signals through (prediction error)
        vs pass (automation).

        Novelty > threshold → target = 1.0 (gate IN, dissolve).
        Novelty ≤ threshold → target = 0.0 (gate OUT, automation).

        The TRN confidence_head outputs a confidence score in (0, 1).
        Low confidence → prediction error → gate IN.
        So we train confidence → (1 − target): high when novel, low when routine.

        Parameters
        ----------
        hidden_at_trn_boundary : (batch, seq, hidden)
            Hidden states at the end of the Layer-2 (subconscious) range.
        """
        novelty = self._compute_novelty(hidden_at_trn_boundary)
        friction_threshold = 1.0 - self.config.trn_gate_target_friction
        gate_target = (novelty > friction_threshold).float()  # 1 = should gate IN

        pooled = hidden_at_trn_boundary.mean(dim=1)  # (batch, hidden)
        confidence = self.ds.trn_gate.confidence_head(pooled).squeeze(-1)  # (batch,)

        # confidence should be LOW when gate_target=1 (prediction error),
        # HIGH when gate_target=0 (automation).
        target_confidence = 1.0 - gate_target
        loss = F.mse_loss(confidence, target_confidence)

        return loss

    # ══════════════════════════════════════════════════════════════════
    # LOSS 2 — Dissolution Compression
    # ══════════════════════════════════════════════════════════════════

    def _dissolution_loss(
        self,
        hidden_at_dissolution: torch.Tensor,
        categories: List[str],
        modulated_logits: torch.Tensor,
        input_ids: torch.Tensor,
    ) -> Tuple[torch.Tensor, Dict[str, float]]:
        """
        The dissolution engine should produce COMPACT, INFORMATION-RICH
        qualia signatures.

        Three sub-losses:
            1. **Reconstruction** — causal LM loss on modulated logits.
               Ensures the dissolution offset preserves enough information
               for coherent output.  Gradient flows through ALL cognitive
               modules (dissolution → metacognitive → spark → fusion →
               modulation → logits).
            2. **Sparsity** — L1 on the raw qualia tensor.  Encourages
               the ``engineered opacity'' compression.
            3. **Contrastive** — same-category prompts should produce
               similar dissolutions; different-category should diverge.

        Returns (total_loss, breakdown_dict).
        """
        device = hidden_at_dissolution.device
        ds = self.ds
        cfg = self.config
        batch_size = hidden_at_dissolution.size(0)

        # ── 2a. Reconstruction: causal LM loss on modulated logits ──
        # Shift logits and labels by one for next-token prediction.
        shift_logits = modulated_logits[:, :-1, :].contiguous()
        shift_labels = input_ids[:, 1:].contiguous()
        recon_loss = F.cross_entropy(
            shift_logits.view(-1, shift_logits.size(-1)),
            shift_labels.view(-1),
            ignore_index=-100,
        )

        # ── 2b. Sparsity: L1 on the raw qualia tensor ──
        pooled = hidden_at_dissolution.mean(dim=1)  # (batch, hidden)
        # Run dissolve_encoder WITH gradients for the L1 term
        raw_qualia = ds.dissolution.dissolve_encoder(pooled)  # (batch, 5)
        gated_qualia = raw_qualia * ds.dissolution.channel_gates.unsqueeze(0)
        sparsity_loss = cfg.dissolution_sparsity_target * F.l1_loss(
            gated_qualia, torch.zeros_like(gated_qualia)
        )

        # ── 2c. Contrastive: same-category similar, different diverge ──
        contrastive_loss = torch.tensor(0.0, device=device)
        n_pairs = 0
        cat_ids = [_CATEGORY_ID.get(c, -1) for c in categories]

        if batch_size >= 2:
            qualia_normed = F.normalize(gated_qualia, dim=-1)  # (batch, 5)
            for i in range(batch_size):
                for j in range(i + 1, batch_size):
                    sim = F.cosine_similarity(
                        qualia_normed[i].unsqueeze(0),
                        qualia_normed[j].unsqueeze(0),
                    ).squeeze(0)
                    if cat_ids[i] == cat_ids[j] and cat_ids[i] >= 0:
                        # Same category → maximise similarity
                        contrastive_loss = contrastive_loss + (1.0 - sim)
                    elif cat_ids[i] != cat_ids[j]:
                        # Different category → push below 0.5
                        contrastive_loss = contrastive_loss + F.relu(sim - 0.5)
                    n_pairs += 1

            if n_pairs > 0:
                contrastive_loss = contrastive_loss / n_pairs

        total = recon_loss + 0.1 * sparsity_loss + 0.05 * contrastive_loss

        breakdown = {
            "recon": float(recon_loss.item()),
            "sparsity": float(sparsity_loss.item()),
            "contrastive": float(contrastive_loss.item()),
        }
        return total, breakdown

    # ══════════════════════════════════════════════════════════════════
    # LOSS 3 — BELBIC Reinforcement Update
    # ══════════════════════════════════════════════════════════════════

    def _compute_reward(
        self,
        hijack_count: int,
        nt_state: Optional[Any],
        sensory_input: torch.Tensor,
    ) -> float:
        """
        Compute a scalar reward for the BELBIC update.

        Reward structure (from ATC training spec):
            +0.5  Generation completed without hijack (stable).
            +1.0  Hijack fired AND output was appropriate (effective
                  circuit breaking — high cortisol/adenosine at end is
                  the proxy for ``appropriate'').
            −0.3  Hijack fired but output was inappropriate (false alarm
                  — hijack fired but NT levels didn't justify it).
            +0.2  Dopamine is high at end (positive engagement).
        """
        reward = 0.0

        if hijack_count == 0:
            reward += 0.5  # Stable generation
        else:
            # Check if hijack was justified by NT levels
            if nt_state is not None:
                cortisol = getattr(nt_state, "cortisol", 0.0)
                adenosine = getattr(nt_state, "adenosine", 0.0)
                if cortisol > 0.7 or adenosine > 0.7:
                    reward += 1.0  # Effective circuit breaking
                else:
                    reward -= 0.3  # False alarm
            else:
                reward -= 0.3  # No NT state — assume false alarm

        # Dopamine bonus
        if nt_state is not None:
            dopamine = getattr(nt_state, "dopamine", 0.0)
            if dopamine > 0.3:
                reward += 0.2

        return reward

    def _belbic_update(
        self,
        sensory_input: torch.Tensor,
        reward: float,
    ) -> None:
        """
        BELBIC reinforcement update — NO gradient.

        Uses the existing ``BELBICDualPathway.update()`` method which
        applies a custom Hebbian-like rule to the amygdala (Go) and
        OFC (NoGo) weight matrices.
        """
        with torch.no_grad():
            self.ds.belbic.update(sensory_input, reward)

    # ══════════════════════════════════════════════════════════════════
    # HOOK HELPERS
    # ══════════════════════════════════════════════════════════════════

    def _install_hooks(self) -> None:
        """
        Install forward hooks on transformer layers at the TRN boundary
        and on the dissolution module to capture intermediate hidden states.
        """
        self._captured = {}
        layers = self.ds._get_transformer_layers()
        boundaries = self.ds._get_layer_boundaries()

        # Hook: end of Layer 2 (subconscious) → TRN gate boundary
        l2_start, l2_end = boundaries["layer2_subconscious"]
        trn_layer = layers[min(l2_end - 1, self.ds.num_layers - 1)]

        def _trn_hook(module, inp, out):
            self._captured["trn_boundary_hidden"] = out[0]

        self._hooks.append(trn_layer.register_forward_hook(_trn_hook))

        # Hook: end of Layer 3 first layer → dissolution boundary
        l3_start, _l3_end = boundaries["layer3_qualia"]
        diss_layer = layers[min(l3_start, self.ds.num_layers - 1)]

        def _diss_hook(module, inp, out):
            self._captured["dissolution_boundary_hidden"] = out[0]

        self._hooks.append(diss_layer.register_forward_hook(_diss_hook))

    def _remove_hooks(self) -> None:
        """Remove all forward hooks.  Captures are preserved for loss computation."""
        for h in self._hooks:
            h.remove()
        self._hooks = []

    def _clear_captures(self) -> None:
        """Clear captured intermediate hidden states."""
        self._captured = {}

    # ══════════════════════════════════════════════════════════════════
    # MAIN TRAINING LOOP
    # ══════════════════════════════════════════════════════════════════

    def train(self) -> Dict[str, Any]:
        """
        Run the self-supervised cognitive training loop.

        Returns a summary dict with final losses and training metadata.
        """
        if not TORCH_AVAILABLE:
            raise RuntimeError("PyTorch required")

        ds = self.ds
        tokenizer = self.model.tokenizer
        cfg = self.config

        # Ensure output directory exists
        out_path = Path(cfg.output_dir)
        out_path.mkdir(parents=True, exist_ok=True)

        # Put cognitive modules in train mode (base model stays eval)
        ds.train()
        self.model.base_model.eval()

        logger.info("Starting ATC cognitive training: %d steps", cfg.max_training_steps)
        train_start = time.time()

        for step in range(cfg.max_training_steps):
            self._step = step
            self.optimizer.zero_grad()

            accum_loss = torch.tensor(0.0, device=self.device)
            step_metrics: Dict[str, float] = {}

            for _accum in range(cfg.gradient_accumulation_steps):
                # ── Sample a batch of prompts ──
                batch_prompts = self._sample_batch()
                categories = [p["category"] for p in batch_prompts]

                # Tokenize (left-pad to equal length)
                input_ids, attention_mask = self._tokenize_batch(
                    batch_prompts, tokenizer
                )

                # ── Reset neurotransmitter shunt for this step ──
                if self.model.nt_shunt is not None:
                    self.model.nt_shunt.reset()

                # ── Install hooks to capture intermediate hidden states ──
                self._install_hooks()

                # ── Forward pass (base model frozen, cognitive modules with grad) ──
                # The forward walks ALL 24 transformer layers, running TRN gate,
                # dissolution, BELBIC, metacognitive loop, and irrational spark
                # at the appropriate layer boundaries.
                try:
                    modulated_logits = ds(
                        input_ids, attention_mask=attention_mask
                    )
                except RuntimeError as exc:
                    self._remove_hooks()
                    if "Ethical veto" in str(exc):
                        logger.warning(
                            "Ethical veto at step %d — skipping batch.", step
                        )
                        continue
                    raise

                # Hooks no longer needed — remove them but keep captures.
                self._remove_hooks()

                # ── Update running mean for novelty estimation ──
                if "trn_boundary_hidden" in self._captured:
                    self._update_running_mean(
                        self._captured["trn_boundary_hidden"]
                    )

                # ── LOSS 1: TRN Gate Calibration ──
                trn_loss = torch.tensor(0.0, device=self.device)
                if "trn_boundary_hidden" in self._captured:
                    trn_loss = self._trn_gate_loss(
                        self._captured["trn_boundary_hidden"]
                    )

                # ── LOSS 2: Dissolution Compression ──
                diss_hidden = self._captured.get(
                    "dissolution_boundary_hidden",
                    self._captured.get("trn_boundary_hidden", input_ids),
                )
                diss_loss, diss_breakdown = self._dissolution_loss(
                    diss_hidden, categories, modulated_logits, input_ids
                )

                # ── LOSS 3: BELBIC Reinforcement (no gradient) ──
                nt_state = self.model.nt_shunt.get_state() if self.model.nt_shunt else None
                hijack_count = ds._hijack_count

                # Build sensory input from current qualia / hidden states
                if ds._current_qualia is not None:
                    q = ds._current_qualia
                    sensory = torch.tensor([[
                        q.valence, q.arousal,
                        ds._last_prediction_error, q.intensity,
                    ]], dtype=torch.float32, device=self.device)
                else:
                    sensory = torch.zeros(1, 4, device=self.device)

                reward = self._compute_reward(hijack_count, nt_state, sensory)
                self._belbic_update(sensory, reward)

                # ── Total loss (TRN + dissolution; BELBIC is separate) ──
                total_loss = trn_loss + 0.1 * diss_loss
                accum_loss = accum_loss + total_loss

                step_metrics.update({
                    "trn_loss": float(trn_loss.item()),
                    **{f"diss_{k}": v for k, v in diss_breakdown.items()},
                    "belbic_reward": reward,
                    "hijack_count": hijack_count,
                })

                # NT state for logging
                if nt_state is not None:
                    step_metrics["nt_ne"] = round(float(nt_state.norepinephrine), 4)
                    step_metrics["nt_cortisol"] = round(float(nt_state.cortisol), 4)
                    step_metrics["nt_dopamine"] = round(float(nt_state.dopamine), 4)
                    step_metrics["nt_adenosine"] = round(float(nt_state.adenosine), 4)

                # Clear captures for next accumulation step
                self._clear_captures()

            # ── Gradient accumulation: average and step ──
            accum_loss = accum_loss / cfg.gradient_accumulation_steps
            accum_loss.backward()
            torch.nn.utils.clip_grad_norm_(
                self._cognitive_params(), max_norm=1.0
            )
            self.optimizer.step()

            # ── Record history ──
            step_metrics["total_loss"] = float(accum_loss.item())
            self._loss_history.append(step_metrics)

            # ── Logging ──
            if step % cfg.log_interval == 0:
                elapsed = time.time() - train_start
                trn_s = step_metrics.get("trn_loss", 0)
                diss_s = step_metrics.get("diss_recon", 0)
                hj = step_metrics.get("hijack_count", 0)
                logger.info(
                    "step %4d/%d  total=%.4f  trn=%.4f  diss_recon=%.4f  "
                    "reward=%+.2f  hijacks=%d  (%.1fs)",
                    step, cfg.max_training_steps,
                    step_metrics.get("total_loss", 0),
                    trn_s, diss_s,
                    step_metrics.get("belbic_reward", 0),
                    hj, elapsed,
                )

            # ── Checkpoint ──
            if step > 0 and step % cfg.save_interval == 0:
                self._save_checkpoint(step, step_metrics)

        # ── Final checkpoint ──
        self._save_checkpoint(cfg.max_training_steps, step_metrics)

        elapsed = time.time() - train_start
        logger.info(
            "Training complete: %d steps in %.1fs.  Checkpoints → %s",
            cfg.max_training_steps, elapsed, cfg.output_dir,
        )

        return {
            "total_steps": cfg.max_training_steps,
            "elapsed_seconds": elapsed,
            "final_losses": step_metrics,
            "output_dir": cfg.output_dir,
        }

    # ── Batch sampling ─────────────────────────────────────────────────

    def _sample_batch(self) -> List[Dict[str, str]]:
        """Sample *batch_size* prompts from the self-supervised pool."""
        import random
        return random.sample(
            SELF_SUPERVISED_PROMPTS, min(self.config.batch_size, len(SELF_SUPERVISED_PROMPTS))
        )

    def _tokenize_batch(
        self,
        prompts: List[Dict[str, str]],
        tokenizer: Any,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """Tokenize a list of prompt dicts to (input_ids, attention_mask)."""
        texts = [p["text"] for p in prompts]
        encoded = tokenizer(
            texts,
            return_tensors="pt",
            padding=True,
            truncation=True,
            max_length=128,
        )
        return (
            encoded["input_ids"].to(self.device),
            encoded["attention_mask"].to(self.device),
        )

    # ── Checkpointing ──────────────────────────────────────────────────

    def _save_checkpoint(
        self, step: int, metrics: Dict[str, float]
    ) -> str:
        """
        Save ONLY the cognitive module state dict (not the base model).

        Returns the path to the saved checkpoint.
        """
        out_path = Path(self.config.output_dir)
        out_path.mkdir(parents=True, exist_ok=True)

        ckpt_path = out_path / f"cognitive_step_{step:06d}.pt"
        checkpoint = {
            "step": step,
            "cognitive_state_dict": self.ds.state_dict(),
            "config": {
                "cognitive_lr": self.config.cognitive_lr,
                "trn_gate_target_friction": self.config.trn_gate_target_friction,
                "dissolution_sparsity_target": self.config.dissolution_sparsity_target,
                "belbic_reward_decay": self.config.belbic_reward_decay,
            },
            "metrics": metrics,
            "timestamp": time.time(),
        }
        torch.save(checkpoint, str(ckpt_path))
        logger.info("Checkpoint saved: %s", ckpt_path)
        return str(ckpt_path)

    @staticmethod
    def load_cognitive_weights(path: str, model: Any) -> None:
        """
        Load previously saved cognitive weights into *model.deep_surgery*.

        Parameters
        ----------
        path : str
            Path to a ``.pt`` checkpoint produced by ``_save_checkpoint``.
        model : NimaModel
            The NIMA unified model (base model weights untouched).
        """
        if not TORCH_AVAILABLE:
            raise RuntimeError("PyTorch required")
        if model.deep_surgery is None:
            raise ValueError("model.deep_surgery is None")

        checkpoint = torch.load(path, map_location="cpu", weights_only=False)
        model.deep_surgery.load_state_dict(checkpoint["cognitive_state_dict"])
        logger.info(
            "Loaded cognitive weights from %s (step %d)",
            path, checkpoint.get("step", "?"),
        )

    # ══════════════════════════════════════════════════════════════════
    # DIAGNOSTICS
    # ══════════════════════════════════════════════════════════════════

    @staticmethod
    def diagnose(model: Any, prompt: str) -> Dict[str, Any]:
        """
        Run a single forward pass and return detailed diagnostics.

        Returns a dict with:
            * Per-module loss breakdown (TRN, dissolution)
            * Neurotransmitter state at each layer boundary
            * Gate decisions, dissolution outcomes, BELBIC gains
            * Recommended adjustments
        """
        if not TORCH_AVAILABLE:
            return {"error": "PyTorch required"}

        ds = model.deep_surgery
        tokenizer = model.tokenizer
        device = next(ds.parameters()).device

        if ds is None:
            return {"error": "deep_surgery not enabled"}

        # ── Run forward pass (eval mode, no grad) ──
        ds.eval()
        model.base_model.eval()
        if model.nt_shunt is not None:
            model.nt_shunt.reset()

        inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=128)
        input_ids = inputs["input_ids"].to(device)
        attention_mask = inputs.get("attention_mask")
        if attention_mask is not None:
            attention_mask = attention_mask.to(device)

        with torch.no_grad():
            try:
                _logits = ds(input_ids, attention_mask=attention_mask)
            except RuntimeError as exc:
                if "Ethical veto" in str(exc):
                    return {"error": f"Ethical veto triggered: {exc}"}
                raise

        metrics = ds.get_consciousness_metrics()
        nt_state = model.nt_shunt.get_state() if model.nt_shunt else None
        audit = ds.get_audit_log()

        # ── Per-module breakdown ──
        report: Dict[str, Any] = {
            "prompt": prompt,
            "neurotransmitters": nt_state.to_dict() if nt_state else {},
            "gate_decisions": [],
            "dissolution_outcomes": [],
            "belbic_gain": metrics.get("belbic_gain", 1.0),
            "metacognitive": {
                "has_qualia": metrics.get("has_qualia", False),
                "hijack_count": metrics.get("hijack_count", 0),
                "dissolutions_fired": metrics.get("dissolutions_fired", 0),
                "dissolutions_deferred": metrics.get("dissolutions_deferred", 0),
            },
            "episodic": {
                "episodes_stored": metrics.get("episodes_stored", 0),
                "reconsolidations": metrics.get("reconsolidations", 0),
                "last_prediction_error": metrics.get("last_prediction_error", 0.0),
            },
            "audit_events": audit[-10:],  # Last 10 events
        }

        # ── Extract gate decisions and dissolution outcomes from audit ──
        for event in audit:
            etype = event.get("event", "")
            if etype == "layer2_subconscious":
                report["gate_decisions"].append({
                    "layer": event.get("layer"),
                    "confidence": event.get("confidence"),
                    "gate_in": event.get("gate_in"),
                    "friction": event.get("friction"),
                })
            elif etype == "layer3_dissolution":
                report["dissolution_outcomes"].append({
                    "layer": event.get("layer"),
                    "valence": event.get("valence"),
                    "arousal": event.get("arousal"),
                    "friction": event.get("friction"),
                })
            elif etype == "amygdala_hijack":
                report.setdefault("hijack_events", []).append({
                    "layer": event.get("layer"),
                    "reason": event.get("reason"),
                })

        # ── Recommended adjustments ──
        adjustments = []
        if nt_state is not None:
            if nt_state.cortisol > 0.8:
                adjustments.append(
                    "Cortisol elevated — consider increasing trn_gate_target_friction "
                    "to let more signals pass subconsciously."
                )
            if nt_state.adenosine > 0.8:
                adjustments.append(
                    "Adenosine high — metacognitive loop may be burning too much ATP. "
                    "Consider raising metacognitive_efficiency_target."
                )
            if nt_state.dopamine < 0.1:
                adjustments.append(
                    "Dopamine low — the system is not finding rewarding patterns. "
                    "Ensure the prompt pool has sufficient variety."
                )
        if metrics.get("dissolutions_deferred", 0) > metrics.get("dissolutions_fired", 0) * 3:
            adjustments.append(
                "Dissolution heavily deferred (alpha phase). "
                "Consider adjusting alpha_freq_hz or alpha_duty_cycle in DissolutionModule."
            )
        if metrics.get("hijack_count", 0) > 2:
            adjustments.append(
                "Multiple hijacks — the irrational spark threshold may be too low. "
                "Review NeurotransmitterState CRITICAL_THRESHOLD."
            )
        if not adjustments:
            adjustments.append("All systems nominal — no adjustments recommended.")

        report["recommended_adjustments"] = adjustments

        return report