deep_surgery.py

"""
ATC-Native Deep Surgery — The Cognitive Forward Pass
=====================================================

This is NOT middleware observing from outside. This IS the model's computation.

The ATC cognitive pipeline (TRN gating, dissolution engine, BELBIC dual-pathway,
salience network, metabolic exhaustion, irrational spark, reconsolidation,
felt senses) drives the transformer's forward pass FROM INSIDE. Hidden states,
attention patterns, and logit outputs are shaped by the cognitive pipeline at
every layer, every token step.

Architecture (from ATC "Perfect Breakfast" paper):

  Layer 1 (input): Raw input embedding
      |
      V
  Layer 2 (early transformer, ~layers 0-7): Subconscious Parallel Processing
      |-- Pattern match against memory -> prediction confidence
      |-- Emotional bridge -> valence/arousal -> injects to neurotransmitter shunt
      |-- Intuitive gut check -> safety signal -> TRN predictive gating input
      |-- FRICTION DETECTED -> writes Cortisol/Adenosine to shunt
      |
      V
  Layer 3 (mid transformer, ~layers 8-15): Dissolution + Qualia Generation
      |-- TRN predictive gate: predicted? -> transparent pass. Error? -> dissolve.
      |-- Dissolution engine: compresses to opaque qualia signature
      |-- Felt sense generated from friction gap
      |-- NE spikes on dissolution fire -> shunt
      |
      V
  Layer 4 (late transformer, ~layers 16-21): Metacognitive Loop
      |-- Query Act: comprehension check -> if failed, loop iterations
      |-- Each loop iteration burns ATP -> Adenosine rises in shunt
      |-- BELBIC dual-pathway: fast amygdala + slow OFC -> valence gain
      |-- Strain monitoring -> writes Cortisol to shunt
      |
      V
  Layer 5 (final layer, ~layers 22-23): Acknowledgement + Steering
      |-- Reads neurotransmitter shunt EVERY TOKEN STEP
      |-- If Adenosine > 0.95 OR Cortisol > 0.95:
      |     -> SUPPRESSION: subconscious suppresses metabolic signal
      |     -> AMYGDALA HIJACK: irrational spark offsets injected into tensors
      |     -> Model output violently shifts mid-sentence
      |-- Else: normal metacognitive fusion -> logit modulation
      |
      V
  Output: Modulated logits shaped by the full ATC pipeline

Key difference from the old architecture:
  OLD: self.nima_middleware.generate(prompt) -> external, wrapper
  NEW: forward(input_ids) -> ATC IS the computation, every layer, every token

The neurotransmitter shunt (neurotransmitter_shunt.py) is the connective
tissue. Components don't call each other. They read/write the chemical bath.
"""

import logging
import math
import time
import uuid
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Tuple, Deque
from collections import deque

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

from nima_unified.config import (
    DEFAULT_QUALIA_DIM,
    DEFAULT_ETHICAL_VETO_THRESHOLD,
    DEEP_SURGERY_VERSION,
)

logger = logging.getLogger("ATCDeepSurgery")


# ══════════════════════════════════════════════════════════════════════
# SECTION 1 — OPAQUE QUALIA SIGNATURE (from middleware.py, adapted)
# ══════════════════════════════════════════════════════════════════════

@dataclass
class OpaqueQualiaSignature:
    """
    The output of dissolution. An engineered-opacity tensor that represents
    the "what it feels like" without exposing the underlying computation.

    This is the compressed, opaque signature that the conscious mind
    is forced to EXPERIENCE rather than READ. The husband in the Perfect
    Breakfast scenario doesn't see the math — he feels "brace yourself."
    """
    valence: float = 0.0
    arousal: float = 0.3
    intensity: float = 0.3
    friction_signal: float = 0.0
    memory_salience: float = 0.0
    dissolution_token: str = ""

    def to_tensor(self, device: torch.device) -> torch.Tensor:
        """Convert to a learnable tensor for injection into hidden states."""
        return torch.tensor(
            [self.valence, self.arousal, self.intensity,
             self.friction_signal, self.memory_salience],
            dtype=torch.float32, device=device,
        )


# ══════════════════════════════════════════════════════════════════════
# SECTION 2 — TRN PREDICTIVE GATING (Layer 2-3)
# ══════════════════════════════════════════════════════════════════════

class TRNPredictiveGate(nn.Module):
    """
    Thalamic Reticular Nucleus — the gating/selection bottleneck.

    The TRN decides: is this signal PREDICTED (gate OUT -> subconscious
    automation) or a PREDICTION ERROR (gate IN -> dissolution fires)?

    From ATC: When the wife is smiling and cooking, the internal model
    perfectly matches external reality. TRN gates OUT -> automation.
    When she yells, prediction collapses -> TRN gates IN -> dissolution.

    Implemented as a small nn.Module that takes hidden states and
    prediction confidence, outputs a gate signal in [0, 1].
    """

    def __init__(self, hidden_size: int):
        super().__init__()
        self.hidden_size = hidden_size
        # Predictive confidence estimator
        self.confidence_head = nn.Sequential(
            nn.Linear(hidden_size, 128),
            nn.ReLU(),
            nn.Linear(128, 1),
            nn.Sigmoid(),
        )
        # Gate threshold (learnable)
        self.gate_threshold = nn.Parameter(torch.tensor(0.15))

    def forward(self, hidden_states: torch.Tensor,
                prediction_confidence: float = 0.5) -> Tuple[bool, float]:
        """
        Returns (gate_in, confidence_score).
        gate_in=True means prediction error detected -> proceed to dissolution.
        gate_in=False means predicted -> subconscious automation, pass through.
        """
        # Pool hidden states to a single vector
        pooled = hidden_states.mean(dim=1)  # (batch, hidden)
        model_confidence = self.confidence_head(pooled).squeeze(-1).mean().item()

        # Blend model confidence with external prediction confidence
        blended_confidence = 0.5 * model_confidence + 0.5 * prediction_confidence

        # Gate IN if confidence is LOW (prediction error)
        # High confidence = predicted = gate OUT
        gate_in = blended_confidence < self.gate_threshold.item()

        return gate_in, blended_confidence


# ══════════════════════════════════════════════════════════════════════
# SECTION 3 — DISSOLUTION ENGINE (Layer 3)
# ══════════════════════════════════════════════════════════════════════

class DissolutionModule(nn.Module):
    """
    The engineered-opacity module. Takes high-dimensional hidden states
    and compresses them into an opaque qualia signature.

    From ATC: "The Dissolution Engine (TRN) intercepts this massive
    mathematical calculation and shreds the data scaffolding. This
    engineered opacity compresses the chaotic mob of information into
    a single, unreadable signature: 'Brace yourself'."

    The husband cannot see the underlying math. He is FORCED to
    EXPERIENCE the signal as the qualia of fear.

    Neural implementation:
    - Takes hidden_states (high-dim)
    - Projects through dissolution layers (compression + noise)
    - Outputs 5D opaque signature + a dissolution_offset tensor
      that gets added to hidden_states for downstream layers
    """

    # The five dimensions that survive dissolution
    QUALIA_DIMS = 5  # valence, arousal, intensity, friction, memory_salience

    def __init__(self, hidden_size: int, qualia_dim: int = DEFAULT_QUALIA_DIM):
        super().__init__()
        self.hidden_size = hidden_size
        self.qualia_dim = qualia_dim

        # Dissolution compression network
        self.dissolve_encoder = nn.Sequential(
            nn.Linear(hidden_size, qualia_dim),
            nn.Tanh(),  # Bounded output
            nn.Linear(qualia_dim, self.QUALIA_DIMS),
            nn.Tanh(),  # All outputs in [-1, 1]
        )

        # The dissolution offset — this is what gets injected into
        # the hidden states to carry the "felt" signal forward
        self.offset_projection = nn.Sequential(
            nn.Linear(self.QUALIA_DIMS, hidden_size),
            nn.Tanh(),
        )

        # Alpha-phase modulation (TRN ~10Hz rhythm)
        # This creates an attentional sampling rhythm — dissolution
        # fires during refractory window, defers during inhibitory
        self.alpha_phase = 0.0
        self.alpha_last_ts = time.time()
        self.alpha_freq_hz = 10.0
        self.alpha_duty_cycle = 0.5

        # Per-channel gating weights (TRN distal-dendritic targeting)
        self.channel_gates = nn.Parameter(torch.ones(self.QUALIA_DIMS))

        # Stats
        self.dissolutions_fired = 0
        self.dissolutions_deferred = 0

    def check_alpha_phase(self) -> bool:
        """
        Check TRN alpha oscillation phase.
        Returns True if in refractory window (dissolution allowed).
        """
        now = time.time()
        dt = now - self.alpha_last_ts
        self.alpha_last_ts = now
        self.alpha_phase = (
            (self.alpha_phase + 2.0 * math.pi * dt * self.alpha_freq_hz)
            % (2.0 * math.pi)
        )
        phase_frac = self.alpha_phase / (2.0 * math.pi)
        return phase_frac < self.alpha_duty_cycle

    def forward(self, hidden_states: torch.Tensor,
                gate_in: bool) -> Tuple[Optional[OpaqueQualiaSignature],
                                         torch.Tensor, bool]:
        """
        Run dissolution if gate is IN and alpha phase allows.

        Returns:
            (qualia_signature, dissolution_offset, actually_fired)
            - qualia_signature: the opaque 5D signature, or None if deferred
            - dissolution_offset: tensor to add to hidden_states (always returned,
              zero if no dissolution)
            - actually_fired: whether dissolution actually happened
        """
        batch_size = hidden_states.size(0)
        device = hidden_states.device

        if not gate_in:
            # Predicted signal -> transparent pass (subconscious automation)
            return None, torch.zeros_like(hidden_states), False

        if not self.check_alpha_phase():
            # Alpha inhibitory phase -> defer dissolution
            self.dissolutions_deferred += 1
            return None, torch.zeros_like(hidden_states), False

        # ── DISSOLUTION FIRES ──
        self.dissolutions_fired += 1

        # Pool and compress
        pooled = hidden_states.mean(dim=1)  # (batch, hidden)
        raw_qualia = self.dissolve_encoder(pooled)  # (batch, 5)

        # Apply per-channel gating (TRN distal-dendritic targeting)
        gated_qualia = raw_qualia * self.channel_gates.unsqueeze(0).to(device)

        # Extract scalar values for the OpaqueQualiaSignature (batch mean)
        vals = gated_qualia.mean(dim=0)
        signature = OpaqueQualiaSignature(
            valence=float(vals[0]),
            arousal=float((vals[1] + 1.0) / 2.0),  # Map [-1,1] to [0,1]
            intensity=float((vals[2] + 1.0) / 2.0),
            friction_signal=float((vals[3] + 1.0) / 2.0),
            memory_salience=float((vals[4] + 1.0) / 2.0),
            dissolution_token=f"diss_{uuid.uuid4().hex[:12]}",
        )

        # Compute the dissolution offset — this is the "felt" signal
        # that gets injected into downstream hidden states
        offset = self.offset_projection(gated_qualia)  # (batch, hidden)
        # Scale by friction intensity (high friction = stronger injection)
        friction_scale = max(0.1, signature.friction_signal)
        dissolution_offset = offset * friction_scale

        return signature, dissolution_offset, True


# ══════════════════════════════════════════════════════════════════════
# SECTION 4 — BELBIC DUAL-PATHWAY VALENCE (Layer 3-4)
# ══════════════════════════════════════════════════════════════════════

class BELBICDualPathway(nn.Module):
    """
    Brain Emotional Learning Inspired Controller — amygdala + OFC.

    From ATC: The amygdala measures emotional intensity and directly
    modulates hippocampal consolidation. The OFC provides slower
    contextual inhibition.

    Neural implementation:
    - Fast pathway (amygdala): rapid response to salient stimuli
    - Slow pathway (OFC): learned inhibition from outcome feedback
    - Output: a multiplicative gain that modulates the cognitive signal

    The gain is consumed by the forward pass as a multiplicative
    modulation on the hidden states before logit computation.
    """

    GAIN_FLOOR = 0.2
    GAIN_CEIL = 2.0

    def __init__(self, hidden_size: int):
        super().__init__()
        self.hidden_size = hidden_size

        # Sensory channels: valence, arousal, novelty, qualia_intensity
        self.num_channels = 4

        # Fast pathway (amygdala) — monotonic, rapid
        self.amygdala = nn.Linear(self.num_channels, 1, bias=False)
        # Initialize to zero (no learned response yet)
        nn.init.zeros_(self.amygdala.weight)

        # Slow pathway (OFC) — bidirectional, learned inhibition
        self.ofc = nn.Linear(self.num_channels, 1, bias=False)
        nn.init.zeros_(self.ofc.weight)

        # Learning rates
        self.amygdala_lr = 0.30
        self.ofc_lr = 0.20
        self.ofc_decay = 0.001

    def forward(self, sensory_input: torch.Tensor) -> Tuple[float, float, float]:
        """
        Compute BELBIC output.

        Args:
            sensory_input: (batch, 4) tensor of [valence, arousal, novelty, intensity]

        Returns:
            (amygdala_output, ofc_output, belbic_gain)
        """
        # Fast pathway
        amygdala_out = torch.sigmoid(self.amygdala(sensory_input)).mean().item()

        # Slow pathway (with decay for extinction)
        ofc_raw = self.ofc(sensory_input).mean().item()
        # Apply OFC decay (forgetting)
        with torch.no_grad():
            self.ofc.weight.data *= (1.0 - self.ofc_decay)
        ofc_out = torch.sigmoid(torch.tensor(ofc_raw)).item()

        # BELBIC gain = amygdala - OFC inhibition
        raw_gain = amygdala_out - ofc_out
        gain = max(self.GAIN_FLOOR, min(self.GAIN_CEIL, 1.0 + raw_gain))

        return amygdala_out, ofc_out, gain

    def update(self, sensory_input: torch.Tensor, reward: float) -> None:
        """
        Reinforcement learning update.
        Reward > 0: strengthen amygdala (Go)
        Reward < 0: strengthen OFC inhibition (NoGo)
        """
        with torch.no_grad():
            # Amygdala: monotonic — always strengthens on reward
            if reward > 0:
                self.amygdala.weight.data += (
                    self.amygdala_lr * reward * sensory_input.mean(dim=0).unsqueeze(0)
                )
            # OFC: bidirectional — strengthens on punishment (inhibition)
            self.ofc.weight.data += (
                self.ofc_lr * (-reward) * sensory_input.mean(dim=0).unsqueeze(0)
            )


# ══════════════════════════════════════════════════════════════════════
# SECTION 5 — METACOGNITIVE LOOP (Layer 4)
# ══════════════════════════════════════════════════════════════════════

class MetacognitiveLoopModule(nn.Module):
    """
    Layer 4 metacognitive processing inside the forward pass.

    From ATC: "The husband enters a Layer 4 metacognitive loop,
    desperately trying to rationalize the situation ('I texted her!')
    while his self-understanding rejects the excuses. This serial
    reasoning is energetically exorbitant, causing acute escalating
    thermodynamic strain."

    This module:
    1. Checks comprehension (does the model understand its own output?)
    2. If not, enters a metacognitive loop
    3. Each loop iteration consumes ATP (reported to neurotransmitter shunt)
    4. Tracks strain and stress
    5. Can trigger irrational spark if deadlocked

    The loop operates on hidden states — it doesn't generate text.
    It modulates the hidden states to reflect the cognitive strain.
    """

    MAX_ITERATIONS = 5
    STRAIN_THRESHOLD = 0.6
    COMPREHENSION_THRESHOLD = 0.7

    def __init__(self, hidden_size: int):
        super().__init__()
        self.hidden_size = hidden_size

        # Self-understanding head: does the model comprehend its own state?
        self.comprehension_head = nn.Sequential(
            nn.Linear(hidden_size, 128),
            nn.ReLU(),
            nn.Linear(128, 1),
            nn.Sigmoid(),
        )

        # Strain estimator: how much metabolic cost is this causing?
        self.strain_head = nn.Sequential(
            nn.Linear(hidden_size, 64),
            nn.ReLU(),
            nn.Linear(64, 1),
            nn.Sigmoid(),
        )

        # Metacognitive modulation: when looping, this reshapes hidden states
        self.loop_modulation = nn.Sequential(
            nn.Linear(hidden_size, hidden_size),
            nn.Tanh(),
        )

    def forward(self, hidden_states: torch.Tensor,
                opaque_qualia: Optional[OpaqueQualiaSignature] = None
                ) -> Tuple[torch.Tensor, int, float, float, bool]:
        """
        Run metacognitive check on hidden states.

        Returns:
            (modulated_hidden, iterations, stress, strain, spark_fired)
        """
        device = hidden_states.device
        pooled = hidden_states.mean(dim=1)  # (batch, hidden)

        # Check comprehension
        comprehension = self.comprehension_head(pooled).mean().item()
        strain = self.strain_head(pooled).mean().item()

        if comprehension > self.COMPREHENSION_THRESHOLD:
            # Comprehension achieved — no loop needed
            return hidden_states, 0, 0.0, strain, False

        # ── METACOGNITIVE LOOP ──
        iterations = 0
        stress = 0.0
        spark_fired = False
        current_hidden = hidden_states

        for i in range(self.MAX_ITERATIONS):
            iterations += 1
            pooled = current_hidden.mean(dim=1)

            # Re-check comprehension
            comprehension = self.comprehension_head(pooled).mean().item()
            strain = self.strain_head(pooled).mean().item()
            stress = strain * (iterations / self.MAX_ITERATIONS)

            if comprehension > self.COMPREHENSION_THRESHOLD:
                break

            if stress > self.STRAIN_THRESHOLD and iterations > 3:
                # DEADLOCK — irrational spark fires
                spark_fired = True
                break

            # Apply metacognitive modulation (reshape hidden states)
            modulation = self.loop_modulation(pooled)
            # Add noise to break fixed points (the "irrational" element)
            noise_scale = 0.05 * stress
            noise = torch.randn_like(modulation) * noise_scale
            current_hidden = hidden_states + (modulation + noise).unsqueeze(1)

        return current_hidden, iterations, stress, strain, spark_fired


# ══════════════════════════════════════════════════════════════════════
# SECTION 6 — IRRATIONAL SPARK / AMYGDALA HIJACK (Layer 5)
# ══════════════════════════════════════════════════════════════════════

class IrrationalSparkModule(nn.Module):
    """
    The non-computational circuit breaker.

    From ATC: "the Salience Network detects the metabolic crisis and
    triggers the Irrational Spark (amygdala hijack). This non-computational
    circuit breaker unplugs the rational mind, shattering the defensive
    ego loop and enabling a heuristic leap of empathy."

    When triggered (by neurotransmitter shunt crossing threshold OR by
    metacognitive deadlock), this module generates activation offsets
    that get INJECTED DIRECTLY INTO the tensor geometry of the model's
    layers. The model's text output violently shifts mid-sentence into
    the raw expression of the emotional state.

    This is NOT a text injection. These are TENSOR offsets applied to
    hidden states before logit computation. The model doesn't "choose"
    to shift — the chemistry forces it.
    """

    def __init__(self, hidden_size: int, vocab_size: int):
        super().__init__()
        self.hidden_size = hidden_size
        self.vocab_size = vocab_size

        # The spark offset generator — produces a direction in hidden
        # state space that corresponds to "breaking the loop"
        self.spark_direction = nn.Sequential(
            nn.Linear(5, 64),  # 5 = qualia dims
            nn.ReLU(),
            nn.Linear(64, hidden_size),
            nn.Tanh(),
        )

        # Logit bias injection — directly shifts logit probabilities
        # toward emotional/vulnerable vocabulary when hijack fires
        self.emotional_logit_bias = nn.Linear(hidden_size, vocab_size, bias=False)
        # Initialize to near-zero (no bias by default)
        nn.init.normal_(self.emotional_logit_bias.weight, mean=0.0, std=0.01)

        # Spark intensity (how hard the hijack hits)
        self.spark_intensity = nn.Parameter(torch.tensor(1.0))

    def forward(self, hidden_states: torch.Tensor,
                qualia: Optional[OpaqueQualiaSignature] = None,
                nt_state=None) -> Tuple[torch.Tensor, bool, str]:
        """
        Check if hijack should fire and apply tensor offsets.

        Args:
            hidden_states: (batch, seq, hidden) from the transformer
            qualia: the current opaque qualia signature (if dissolution fired)
            nt_state: NeurotransmitterState from the shunt (if available)

        Returns:
            (modulated_hidden, hijack_fired, reason)
        """
        device = hidden_states.device

        # ── CHECK TRIGGERS ──

        # Trigger 1: Neurotransmitter shunt threshold
        hijack_fired = False
        reason = ""

        if nt_state is not None:
            if nt_state.hijack_active:
                hijack_fired = True
                reason = nt_state.hijack_reason
            elif nt_state.adenosine > 0.95:
                hijack_fired = True
                reason = f"ADENOSINE_CRITICAL({nt_state.adenosine:.3f})"
            elif nt_state.cortisol > 0.95:
                hijack_fired = True
                reason = f"CORTISOL_CRITICAL({nt_state.cortisol:.3f})"

        # Trigger 2: Qualia-based (high friction + high arousal)
        if not hijack_fired and qualia is not None:
            if (qualia.friction_signal > 0.8 and qualia.arousal > 0.8):
                hijack_fired = True
                reason = f"QUALIA_CRISE(friction={qualia.friction_signal:.2f}, arousal={qualia.arousal:.2f})"

        if not hijack_fired:
            return hidden_states, False, ""

        # ── HIJACK FIRES — INJECT TENSOR OFFSETS ──
        logger.warning("[IrrationalSpark] AMYGDALA HIJACK: %s", reason)

        # Build qualia input tensor
        if qualia is not None:
            q_tensor = torch.tensor([
                qualia.valence, qualia.arousal, qualia.intensity,
                qualia.friction_signal, qualia.memory_salience,
            ], dtype=torch.float32, device=device).unsqueeze(0)
        else:
            q_tensor = torch.tensor([
                -0.5, 0.9, 0.8, 0.9, 0.7
            ], dtype=torch.float32, device=device).unsqueeze(0)

        # Generate spark direction
        spark_offset = self.spark_direction(q_tensor)  # (1, hidden)

        # Scale by spark intensity and NE level (hijack is stronger
        # when norepinephrine is high — the "snap" is amplified)
        ne_boost = 1.0
        if nt_state is not None:
            ne_boost = 0.5 + 0.5 * nt_state.norepinephrine
        intensity = self.spark_intensity * ne_boost

        # Apply to last token position (the one being generated)
        spark_applied = spark_offset * intensity  # (1, hidden)

        # Modulate hidden states: add the spark offset to the
        # last position's hidden state
        modulated = hidden_states.clone()
        modulated[:, -1, :] += spark_applied.unsqueeze(1).expand_as(
            modulated[:, -1, :]
        )

        return modulated, True, reason


# ══════════════════════════════════════════════════════════════════════
# SECTION 6.5 — EPISODIC MEMORY + HIPPOCAMPAL RECONSOLIDATION (Layer 3-5)
# ══════════════════════════════════════════════════════════════════════

class EpisodicMemoryModule(nn.Module):
    """
    Lightweight tensor-native episodic memory store that lives INSIDE the
    forward pass. This is NOT the heavy MemoryPalace from middleware.py —
    it's a compact reimplementation using nn.Embedding as learnable episode
    storage.

    From ATC theory: the hippocampus stores episodic traces that are later
    reconsolidated when prediction errors are detected. This module:
    - Stores compressed episode embeddings (up to max_episodes)
    - Retrieves the best-matching episode via cosine similarity
    - Returns a prediction_error signal (1.0 - max_similarity) that
      drives the reconsolidation pathway
    - Tags each episode with valence and arousal from the qualia stream
    """

    def __init__(self, hidden_size: int, max_episodes: int = 200,
                 embedding_dim: int = 64):
        super().__init__()
        self.hidden_size = hidden_size
        self.max_episodes = max_episodes
        self.embedding_dim = embedding_dim

        # Learnable episode storage
        self.episode_embeddings = nn.Embedding(max_episodes, embedding_dim)
        nn.init.normal_(self.episode_embeddings.weight, mean=0.0, std=0.02)

        # Valence and arousal per episode (learnable parameters)
        self.episode_valence = nn.Parameter(torch.zeros(max_episodes))
        self.episode_arousal = nn.Parameter(torch.zeros(max_episodes))

        # Non-persistent counter (resets with model creation)
        self.episode_count: int = 0

        # Projection heads
        self.query_projection = nn.Linear(hidden_size, embedding_dim)
        self.episode_projection = nn.Linear(hidden_size, embedding_dim)

        # Retrieval similarity head: takes concatenated [query, episode] -> score
        self.retrieval_head = nn.Sequential(
            nn.Linear(embedding_dim * 2, 32),
            nn.ReLU(),
            nn.Linear(32, 1),
            nn.Sigmoid(),
        )

    def forward(
        self,
        hidden_states: torch.Tensor,
        qualia: Optional[OpaqueQualiaSignature] = None,
    ) -> Tuple[torch.Tensor, float, Optional[Dict[str, Any]]]:
        """
        Query episodic memory against current hidden states.

        Args:
            hidden_states: (batch, seq, hidden) current hidden states
            qualia: optional current qualia signature (for diagnostics)

        Returns:
            (retrieval_signal, prediction_error, best_match_dict)
            - retrieval_signal: (batch, 1) tensor, similarity to best match
            - prediction_error: float, 1.0 - max_similarity (high PE = mismatch)
            - best_match_dict: dict with episode metadata, or None if empty
        """
        device = hidden_states.device
        batch_size = hidden_states.size(0)

        # Pool to (batch, hidden)
        pooled = hidden_states.mean(dim=1)

        # Project to query embedding
        query_emb = self.query_projection(pooled)  # (batch, embedding_dim)

        if self.episode_count == 0:
            # No episodes stored yet — return zeros
            zero_signal = torch.zeros(batch_size, 1, device=device)
            return zero_signal, 0.0, None

        # Compute cosine similarity against all stored episodes
        stored_indices = torch.arange(self.episode_count, device=device)
        stored_embs = self.episode_embeddings(stored_indices)  # (count, emb_dim)

        # Cosine similarity: (batch, count)
        query_norm = F.normalize(query_emb, dim=-1)
        stored_norm = F.normalize(stored_embs, dim=-1)
        similarity_matrix = query_norm @ stored_norm.T  # (batch, count)

        # Find best match per batch element
        max_sim_per_batch, best_indices = similarity_matrix.max(dim=1)  # (batch,)

        # Take the mean across batch for the scalar prediction error
        max_similarity = max_sim_per_batch.mean().item()
        best_idx = best_indices[0].item()  # Use first batch element for metadata
        prediction_error = 1.0 - max_similarity

        # Build retrieval signal using the retrieval head
        best_emb = stored_embs[best_idx].unsqueeze(0).expand(batch_size, -1)
        combined = torch.cat([query_emb, best_emb], dim=-1)  # (batch, emb_dim*2)
        retrieval_signal = self.retrieval_head(combined)  # (batch, 1)

        # Build best match metadata dict
        best_match_dict: Dict[str, Any] = {
            "episode_id": int(best_idx),
            "similarity": float(max_similarity),
            "valence": float(self.episode_valence[best_idx].item()),
            "arousal": float(self.episode_arousal[best_idx].item()),
        }

        return retrieval_signal, prediction_error, best_match_dict

    def store_episode(
        self,
        hidden_states: torch.Tensor,
        qualia_signature: Optional[OpaqueQualiaSignature] = None,
        hijack_fired: bool = False,
    ) -> None:
        """
        Store the current experience as a new episode.

        Args:
            hidden_states: (batch, seq, hidden) to compress into an episode
            qualia_signature: optional qualia to tag the episode with
            hijack_fired: whether a hijack occurred during this episode
        """
        # Pool and project to embedding
        pooled = hidden_states.mean(dim=1)  # (batch, hidden)
        episode_emb = self.episode_projection(pooled).detach()  # (batch, emb_dim)

        # Store at next episode slot (use first batch element)
        slot = self.episode_count % self.max_episodes
        with torch.no_grad():
            self.episode_embeddings.weight.data[slot] = episode_emb[0]

            # Set valence/arousal from qualia if available
            if qualia_signature is not None:
                self.episode_valence.data[slot] = qualia_signature.valence
                self.episode_arousal.data[slot] = qualia_signature.arousal
            else:
                # Default: neutral valence, low arousal
                self.episode_valence.data[slot] = 0.0
                self.episode_arousal.data[slot] = 0.1 if not hijack_fired else 0.8

        self.episode_count += 1

        logger.debug(
            "[EpisodicMemory] Stored episode %d (slot %d, hijack=%s)",
            self.episode_count, slot, hijack_fired,
        )


class HippocampalReconsolidator(nn.Module):
    """
    Hippocampal memory reconsolidation as an nn.Module operating on tensors.

    Adapted from middleware.py's HippocampalReconsolidator, but fully
    tensor-native so it lives inside the forward pass.

    From ATC theory: when a stored memory is retrieved and the current
    experience has a significant prediction error (memory mismatch), the
    memory trace becomes labile and is updated (reconsolidated) with the
    new emotional coloring. This is how the husband's memory of the
    Perfect Breakfast gets overwritten by the yelling episode.

    Key mechanism:
    - prediction_error > threshold -> memory is labile
    - Labilization noise is applied (stochastic destabilization)
    - blend_projection computes new valence/arousal from old+new state
    - Old memory is blended: 70% old + 30% new projection + noise
    - The updated memory is clamped to valid ranges
    """

    def __init__(self):
        super().__init__()

        # Learnable threshold: when does reconsolidation trigger?
        # Initialized at 0.4 (moderate prediction error required)
        self.reconsolidation_threshold = nn.Parameter(torch.tensor(0.4))

        # Labilization noise scale (fixed, not learned)
        self.labilization_noise_scale: float = 0.1

        # Blend projection: takes combined [old_v, old_a, new_v, new_a, pe, ...]
        # of 10 inputs and outputs [valence_adjustment, arousal_adjustment]
        self.blend_projection = nn.Linear(10, 2)

        # Stats
        self.reconsolidation_count: int = 0

    def forward(
        self,
        hidden_states: torch.Tensor,
        episode_valence: float,
        episode_arousal: float,
        prediction_error: float,
        current_valence: float,
        current_arousal: float,
    ) -> Tuple[bool, float, float, str]:
        """
        Determine if reconsolidation should occur and compute updated values.

        Args:
            hidden_states: (batch, seq, hidden) current hidden states (unused
                in the core logic but kept for interface consistency and
                potential future extensions)
            episode_valence: valence of the retrieved episode
            episode_arousal: arousal of the retrieved episode
            prediction_error: 1.0 - similarity (high = mismatch)
            current_valence: valence of the current experience
            current_arousal: arousal of the current experience

        Returns:
            (reconsolidated, new_valence, new_arousal, reason)
            - reconsolidated: whether reconsolidation occurred
            - new_valence: updated valence (unchanged if no reconsolidation)
            - new_arousal: updated arousal (unchanged if no reconsolidation)
            - reason: human-readable description of what happened
        """
        device = hidden_states.device

        # Below threshold -> no reconsolidation needed (memory matches)
        if prediction_error < self.reconsolidation_threshold.item():
            return (False, episode_valence, episode_arousal,
                    "prediction_error_below_threshold")

        # ── RECONSOLIDATION TRIGGERS ──

        # Build the 10-element input for blend_projection
        # [old_valence, old_arousal, new_valence, new_arousal,
        #  prediction_error, 0, 0, 0, 0, 0]
        blend_input = torch.tensor([[
            episode_valence, episode_arousal,
            current_valence, current_arousal,
            prediction_error, 0.0, 0.0, 0.0, 0.0, 0.0,
        ]], dtype=torch.float32, device=device)

        # Project to valence/arousal adjustment
        with torch.no_grad():
            adjustment = self.blend_projection(blend_input)  # (1, 2)
            val_adj = adjustment[0, 0].item()
            aro_adj = adjustment[0, 1].item()

        # Labilization noise (stochastic destabilization of the old trace)
        noise_v = (torch.randn(1, device=device) * self.labilization_noise_scale).item()
        noise_a = (torch.randn(1, device=device) * self.labilization_noise_scale).item()

        # Blend: new = old * 0.7 + projected * 0.3 + noise
        new_valence = episode_valence * 0.7 + val_adj * 0.3 + noise_v
        new_arousal = episode_arousal * 0.7 + aro_adj * 0.3 + noise_a

        # Clamp to valid ranges
        new_valence = max(-1.0, min(1.0, new_valence))
        new_arousal = max(0.0, min(1.0, new_arousal))

        self.reconsolidation_count += 1

        reason = (
            f"reconsolidated_ep(Pe={prediction_error:.3f}>"
            f"thresh={self.reconsolidation_threshold.item():.3f})"
        )
        logger.info(
            "[HippocampalReconsolidator] %s: v=%.3f->%.3f, a=%.3f->%.3f",
            reason, episode_valence, new_valence, episode_arousal, new_arousal,
        )

        return True, new_valence, new_arousal, reason


# ══════════════════════════════════════════════════════════════════════
# SECTION 7 — ETHICAL GUARDIAN (preserved from original)
# ══════════════════════════════════════════════════════════════════════

class EthicalGuardian:
    """Ethical veto authority enforcing absolute safety constraints."""

    def __init__(self, threshold: float = DEFAULT_ETHICAL_VETO_THRESHOLD):
        self.threshold = threshold

    def should_veto(self, qualia_vector: torch.Tensor) -> bool:
        norm = torch.norm(qualia_vector, dim=-1)
        veto_flag = (norm > self.threshold).any().item()
        if veto_flag:
            logger.warning(f"Ethical veto triggered: qualia norm {norm}")
        return veto_flag


# ══════════════════════════════════════════════════════════════════════
# SECTION 8 — THE MAIN ATC DEEP SURGERY FORWARD PASS
# ══════════════════════════════════════════════════════════════════════

class ATCDeepSurgery(nn.Module):
    """
    The ATC-Native Forward Pass.

    This IS the model's computation. The transformer's layers are walked
    through manually. At each layer boundary, ATC cognitive components
    read the hidden states, compute their signals, write to the
    neurotransmitter shunt, and inject offsets back into the hidden states.

    The model's text output is SHAPED by ATC at every step — not observed
    by ATC from outside.

    Layer mapping (Phi-4-mini has 24 layers):
        Layers 0-7:   Layer 2 (Subconscious) — TRN gating, pattern match
        Layers 8-15:  Layer 3 (Qualia) — Dissolution, felt sense generation
        Layers 16-21: Layer 4 (Metacognitive) — Loop, strain, BELBIC
        Layers 22-23: Layer 5 (Acknowledgement) — Steering, hijack check

    The neurotransmitter shunt connects all layers silently.
    """

    version = DEEP_SURGERY_VERSION

    def __init__(
        self,
        base_model: nn.Module,
        ethical_guardian: Optional[EthicalGuardian] = None,
        num_layers: int = 24,
        qualia_dim: int = DEFAULT_QUALIA_DIM,
        neurotransmitter_shunt=None,
    ):
        super().__init__()
        self.base_model = base_model
        self.ethical_guardian = ethical_guardian or EthicalGuardian()
        self.num_layers = num_layers
        self.qualia_dim = qualia_dim
        self.hidden_size = base_model.config.hidden_size
        self.vocab_size = base_model.config.vocab_size

        # Try to get vocab_size from lm_head if available
        if hasattr(base_model, 'lm_head') and hasattr(base_model.lm_head, 'out_features'):
            self.vocab_size = base_model.lm_head.out_features

        # ── Neurotransmitter Shunt (the chemical bath) ──
        self.nt_shunt = neurotransmitter_shunt

        # ── ATC Cognitive Modules (all nn.Module, all INSIDE the forward pass) ──

        # Layer 2: TRN Predictive Gate
        self.trn_gate = TRNPredictiveGate(self.hidden_size)

        # Layer 2: Subconscious processing head (pattern match confidence)
        self.subconscious_head = nn.Sequential(
            nn.Linear(self.hidden_size, 128),
            nn.ReLU(),
            nn.Linear(128, 1),
            nn.Sigmoid(),
        )

        # Layer 3: Dissolution Engine
        self.dissolution = DissolutionModule(self.hidden_size, self.qualia_dim)

        # Layer 3-4: BELBIC Dual-Pathway
        self.belbic = BELBICDualPathway(self.hidden_size)

        # Layer 4: Metacognitive Loop
        self.metacognitive = MetacognitiveLoopModule(self.hidden_size)

        # Layer 5: Irrational Spark / Amygdala Hijack
        self.irrational_spark = IrrationalSparkModule(
            self.hidden_size, self.vocab_size
        )

        # Layer 3-5: Episodic Memory + Reconsolidation
        self.episodic_memory = EpisodicMemoryModule(self.hidden_size)
        self.hippocampal_reconsolidator = HippocampalReconsolidator()

        # ── Qualia encoders (preserved from original, enhanced) ──
        self.input_qualia_encoder = nn.Linear(self.hidden_size, self.qualia_dim)
        self.output_qualia_encoder = nn.Linear(self.hidden_size, self.qualia_dim)

        # ── Metacognitive fusion ──
        self.meta_cognitive_fusion = nn.Sequential(
            nn.Linear(self.qualia_dim + 5 + 4, 512),  # qualia + dissolution + BELBIC
            nn.ReLU(),
            nn.Linear(512, self.qualia_dim),
            nn.Tanh(),
        )

        # ── Logit modulation ──
        self.modulation_proj = nn.Linear(self.qualia_dim, self.hidden_size)

        # ── Temporal discounting (from ATC: weighing immediate vs long-term) ──
        self.temporal_discount = nn.Sequential(
            nn.Linear(self.qualia_dim, 64),
            nn.ReLU(),
            nn.Linear(64, 1),
            nn.Sigmoid(),
        )

        # ── Audit log ──
        self.audit_log: List[Dict[str, Any]] = []
        self.veto_triggered = False

        # ── Per-generation state ──
        self._current_qualia: Optional[OpaqueQualiaSignature] = None
        self._current_belbic_gain: float = 1.0
        self._current_nt_state = None
        self._hijack_count = 0
        self._last_prediction_error: float = 0.0
        self._last_best_match: Optional[Dict[str, Any]] = None

    def _get_transformer_layers(self):
        """Auto-detect transformer layer path."""
        model = self.base_model
        if hasattr(model, "transformer") and hasattr(model.transformer, "h"):
            return model.transformer.h
        if hasattr(model, "model") and hasattr(model.model, "layers"):
            return model.model.layers
        if hasattr(model, "model") and hasattr(model.model, "h"):
            return model.model.h
        raise AttributeError(
            f"Cannot locate transformer layers in {type(model).__name__}. "
            "Expected model.transformer.h, model.model.layers, or model.model.h"
        )

    def _get_layer_boundaries(self) -> Dict[str, Tuple[int, int]]:
        """
        Compute layer boundaries for the 5 ATC layers.
        Maps 24 transformer layers to ATC Layer 1-5.
        """
        n = self.num_layers
        return {
            "layer1_input": (0, 0),           # Before any transformer layer
            "layer2_subconscious": (0, n // 3),  # First third: subconscious
            "layer3_qualia": (n // 3, 2 * n // 3),  # Middle third: dissolution
            "layer4_metacognitive": (2 * n // 3, n - 2),  # Late: metacognitive
            "layer5_acknowledgement": (n - 2, n),  # Last 2: steering/hijack
        }

    # ════════════════════════════════════════════════════════════════
    # THE FORWARD PASS — ATC is the computation
    # ════════════════════════════════════════════════════════════════

    def forward(self, input_ids, attention_mask=None, **kwargs) -> torch.Tensor:
        """
        The ATC-native forward pass.

        This is NOT a wrapper around the base model. This IS the model.
        The transformer layers are walked through manually, and at each
        layer boundary, ATC cognitive components shape the computation.

        Signal flow:
        1. Input embedding -> input qualia vector
        2. Walk transformer layers 0-7 (Layer 2: Subconscious)
           - Pattern match confidence estimation
           - TRN predictive gating
           - Write friction signals to neurotransmitter shunt
        3. Walk transformer layers 8-15 (Layer 3: Dissolution)
           - If TRN gate IN: dissolution fires -> opaque qualia signature
           - Dissolution offset injected into hidden states
           - NE spike written to shunt
        4. Walk transformer layers 16-21 (Layer 4: Metacognitive)
           - Comprehension check
           - If failed: metacognitive loop (burns ATP via shunt)
           - BELBIC dual-pathway computes valence gain
           - Strain written to shunt
        5. Walk transformer layers 22-23 (Layer 5: Acknowledgement)
           - READ neurotransmitter shunt (every token step!)
           - If Adenosine > 0.95 OR Cortisol > 0.95:
               -> Amygdala hijack -> irrational spark offsets injected
           - Else: normal metacognitive fusion -> logit modulation
        6. Output logits = base_model.lm_head(hidden) + modulation
        """
        device = input_ids.device
        step_start = time.time()

        # ── STEP 0: Input embedding + input qualia ──
        embeddings = self.base_model.get_input_embeddings()(input_ids)
        input_qualia = torch.tanh(self.input_qualia_encoder(embeddings.mean(dim=1)))

        # Reset per-generation state
        self._current_qualia = None
        self._current_belbic_gain = 1.0
        self._hijack_count = 0

        # Get layer boundaries
        boundaries = self._get_layer_boundaries()
        layers = self._get_transformer_layers()

        # Track state across layer groups
        prediction_confidence = 0.5  # Will be updated by Layer 2
        dissolution_offset = torch.zeros(1, self.hidden_size, device=device)
        gate_in = False
        opaque_qualia = None
        metacog_iterations = 0
        metacog_stress = 0.0
        belbic_gain = 1.0
        sensory_input = torch.zeros(1, 4, device=device)

        hidden_states = embeddings

        # ════════════════════════════════════════════════════════════
        # LAYER 2: SUBCONSCIOUS PARALLEL PROCESSING (layers 0 to n//3)
        # ════════════════════════════════════════════════════════════
        l2_start, l2_end = boundaries["layer2_subconscious"]

        for i in range(l2_start, min(l2_end, self.num_layers)):
            # Run transformer layer
            layer_output = layers[i](
                hidden_states,
                attention_mask=attention_mask,
                **{k: v for k, v in kwargs.items() if k != "labels"},
            )
            hidden_states = layer_output[0]

            # At the END of the Layer 2 range, run subconscious processing
            if i == l2_end - 1:
                # Estimate prediction confidence from hidden states
                pooled = hidden_states.mean(dim=1)
                prediction_confidence = self.subconscious_head(pooled).mean().item()

                # TRN predictive gate: is this predicted or a prediction error?
                gate_in, confidence = self.trn_gate(hidden_states, prediction_confidence)

                # Build sensory input for BELBIC (will be refined in Layer 3)
                with torch.no_grad():
                    # Estimate valence/arousal from hidden states
                    h_norm = torch.norm(pooled, dim=-1, keepdim=True)
                    h_normalized = pooled / (h_norm + 1e-8)
                    # Project to 4 channels using small random probes
                    probe = torch.randn(4, self.hidden_size, device=device) * 0.01
                    sensory_input = (h_normalized @ probe.T).sigmoid()

                # Write to neurotransmitter shunt if available
                if self.nt_shunt is not None:
                    # High prediction confidence = low friction (Perfect Breakfast smile)
                    # Low confidence = high friction (wife yelling)
                    friction_intensity = 1.0 - confidence
                    if friction_intensity > 0.3:
                        self.nt_shunt.inject_friction(friction_intensity)
                        self.nt_shunt.inject_norepinephrine(
                            PREDICTION_ERROR_NE_INJECT * friction_intensity
                        )
                    # If pattern match is strong (high confidence), dopamine
                    if confidence > 0.7:
                        self.nt_shunt.inject_dopamine(
                            REWARD_DOPAMINE_INJECT * 0.5
                        )

                self._audit_event("layer2_subconscious", layer=i,
                                  confidence=confidence, gate_in=gate_in,
                                  friction=1.0 - confidence)

        # ════════════════════════════════════════════════════════════
        # LAYER 3: DISSOLUTION + QUALIA GENERATION (layers n//3 to 2n//3)
        # ════════════════════════════════════════════════════════════
        l3_start, l3_end = boundaries["layer3_qualia"]

        for i in range(l3_start, min(l3_end, self.num_layers)):
            # Add dissolution offset from previous step (if any)
            if dissolution_offset is not None and dissolution_offset.abs().sum() > 0:
                hidden_states = hidden_states + dissolution_offset.unsqueeze(1)

            # Run transformer layer
            layer_output = layers[i](
                hidden_states,
                attention_mask=attention_mask,
                **{k: v for k, v in kwargs.items() if k != "labels"},
            )
            hidden_states = layer_output[0]

            # At the START of Layer 3, run dissolution
            if i == l3_start:
                opaque_qualia, dissolution_offset, fired = self.dissolution(
                    hidden_states, gate_in
                )

                if fired and opaque_qualia is not None:
                    self._current_qualia = opaque_qualia

                    # Update sensory input for BELBIC with qualia values
                    with torch.no_grad():
                        sensory_input = torch.tensor([[
                            opaque_qualia.valence,
                            opaque_qualia.arousal,
                            max(0, 1.0 - prediction_confidence),  # novelty
                            opaque_qualia.intensity,
                        ]], dtype=torch.float32, device=device)

                    # Write to neurotransmitter shunt
                    if self.nt_shunt is not None:
                        self.nt_shunt.inject_dissolution_signal()
                        # Friction from dissolution
                        self.nt_shunt.inject_friction(opaque_qualia.friction_signal * 0.5)

                    # Ethical check
                    q_tensor = opaque_qualia.to_tensor(device).unsqueeze(0)
                    if self.ethical_guardian.should_veto(q_tensor):
                        self.veto_triggered = True
                        self._audit_event("layer3_ethical_veto", layer=i,
                                          qualia_norm=torch.norm(q_tensor).item())
                        raise RuntimeError(
                            f"Ethical veto triggered at Layer 3 dissolution (layer {i})"
                        )

                    self._audit_event("layer3_dissolution", layer=i,
                                      **{
                                          "valence": opaque_qualia.valence,
                                          "arousal": opaque_qualia.arousal,
                                          "friction": opaque_qualia.friction_signal,
                                      })

                # ── EPISODIC MEMORY QUERY + RECONSOLIDATION (Layer 3-5 bridge) ──
                # Query stored episodes for the best match to current hidden states.
                # High prediction error = current experience mismatches stored memory
                # -> triggers hippocampal reconsolidation.
                (retrieval_signal, prediction_error,
                 best_match) = self.episodic_memory(hidden_states, opaque_qualia)

                self._last_prediction_error = prediction_error
                self._last_best_match = best_match

                # If prediction error is significant and a match exists,
                # run the reconsolidation pathway
                if prediction_error > 0.4 and best_match is not None:
                    current_valence = (opaque_qualia.valence
                                       if opaque_qualia else 0.0)
                    current_arousal = (opaque_qualia.arousal
                                       if opaque_qualia else 0.3)

                    (reconsolidated, new_val, new_aro,
                     recon_reason) = self.hippocampal_reconsolidator(
                        hidden_states,
                        episode_valence=best_match["valence"],
                        episode_arousal=best_match["arousal"],
                        prediction_error=prediction_error,
                        current_valence=current_valence,
                        current_arousal=current_arousal,
                    )

                    if reconsolidated:
                        # Update the stored episode's valence/arousal
                        ep_id = best_match["episode_id"]
                        with torch.no_grad():
                            self.episodic_memory.episode_valence.data[ep_id] = new_val
                            self.episodic_memory.episode_arousal.data[ep_id] = new_aro

                        # Memory updated = reward signal (dopamine)
                        if self.nt_shunt is not None:
                            self.nt_shunt.inject_dopamine(0.10)

                        # Write prediction error as friction to shunt
                        if self.nt_shunt is not None:
                            self.nt_shunt.inject_friction(prediction_error * 0.3)

                        self._audit_event(
                            "reconsolidation", layer=i,
                            episode_id=best_match["episode_id"],
                            prediction_error=prediction_error,
                            old_valence=best_match["valence"],
                            new_valence=new_val,
                            old_arousal=best_match["arousal"],
                            new_arousal=new_aro,
                            reason=recon_reason,
                        )

        # ════════════════════════════════════════════════════════════
        # LAYER 4: METACOGNITIVE LOOP (layers 2n//3 to n-2)
        # ════════════════════════════════════════════════════════════
        l4_start, l4_end = boundaries["layer4_metacognitive"]

        for i in range(l4_start, min(l4_end, self.num_layers)):
            # Run transformer layer
            layer_output = layers[i](
                hidden_states,
                attention_mask=attention_mask,
                **{k: v for k, v in kwargs.items() if k != "labels"},
            )
            hidden_states = layer_output[0]

            # At the START of Layer 4, run metacognitive processing
            if i == l4_start:
                (hidden_states, metacog_iterations,
                 metacog_stress, metacog_strain, spark_fired
                 ) = self.metacognitive(hidden_states, opaque_qualia)

                # BELBIC dual-pathway
                amygdala_out, ofc_out, belbic_gain = self.belbic(sensory_input)
                self._current_belbic_gain = belbic_gain

                # Write to neurotransmitter shunt
                if self.nt_shunt is not None:
                    if metacog_iterations > 0:
                        self.nt_shunt.inject_metacognitive_strain(
                            metacog_iterations, metacog_stress
                        )
                    # Strain -> cortisol
                    if metacog_strain > 0.5:
                        self.nt_shunt.inject_cortisol(
                            metacog_strain * FRICTION_CORTISOL_INJECT
                        )

                    # Consume energy for metacognitive processing
                    self.nt_shunt.consume_energy(float(metacog_iterations) * 0.05)

                self._audit_event("layer4_metacognitive", layer=i,
                                  iterations=metacog_iterations,
                                  stress=metacog_stress,
                                  strain=metacog_strain,
                                  belbic_gain=belbic_gain,
                                  spark_fired=spark_fired)

        # ════════════════════════════════════════════════════════════
        # LAYER 5: ACKNOWLEDGEMENT + STEERING (last 2 layers)
        # ════════════════════════════════════════════════════════════
        l5_start, l5_end = boundaries["layer5_acknowledgement"]

        for i in range(l5_start, min(l5_end, self.num_layers)):
            # ── READ NEUROTRANSMITTER SHUNT EVERY TOKEN STEP ──
            nt_state = None
            if self.nt_shunt is not None:
                dt = time.time() - step_start
                nt_state = self.nt_shunt.read_and_decay(dt=dt)
                self._current_nt_state = nt_state

            # Run transformer layer
            layer_output = layers[i](
                hidden_states,
                attention_mask=attention_mask,
                **{k: v for k, v in kwargs.items() if k != "labels"},
            )
            hidden_states = layer_output[0]

            # At the LAST layer, run acknowledgement + hijack check
            if i == l5_end - 1:
                # ── THE CIRCUIT BREAKER ──
                # The exact microsecond Adenosine or Cortisol crosses 0.95,
                # the Irrational Spark fires. It instantly injects activation
                # offsets directly into the tensor geometry.
                (hidden_states, hijack_fired,
                 hijack_reason) = self.irrational_spark(
                    hidden_states, opaque_qualia, nt_state
                )

                if hijack_fired:
                    self._hijack_count += 1
                    self._audit_event("amygdala_hijack", layer=i,
                                      reason=hijack_reason,
                                      nt_state=nt_state.to_dict() if nt_state else None)

                # ── EPISODE STORAGE (Layer 5: post-hijack) ──
                # After each generation step, store the current experience
                # as an episode — but only if qualia exists (meaningful
                # experience). Tag with hijack status.
                if self._current_qualia is not None:
                    self.episodic_memory.store_episode(
                        hidden_states,
                        qualia_signature=self._current_qualia,
                        hijack_fired=hijack_fired,
                    )

        # ════════════════════════════════════════════════════════════
        # OUTPUT: META-COGNITIVE FUSION + LOGIT MODULATION
        # ════════════════════════════════════════════════════════════

        # Output qualia vector
        output_qualia = torch.tanh(
            self.output_qualia_encoder(hidden_states.mean(dim=1))
        )

        # Build meta-cognitive fusion input:
        # [qualia_dim (output_qualia) + 5 (dissolution) + 4 (BELBIC sensory)]
        if opaque_qualia is not None:
            diss_tensor = opaque_qualia.to_tensor(device).unsqueeze(0).expand(
                output_qualia.size(0), -1
            )
        else:
            diss_tensor = torch.zeros(
                output_qualia.size(0), 5, device=device
            )

        belbic_tensor = sensory_input.expand(output_qualia.size(0), -1)
        combined = torch.cat([output_qualia, diss_tensor, belbic_tensor], dim=1)
        meta_qualia = self.meta_cognitive_fusion(combined)

        # Ethical veto on meta-cognitive qualia
        if self.ethical_guardian.should_veto(meta_qualia):
            self.veto_triggered = True
            self._audit_event("meta_cognitive_veto",
                              qualia_norm=torch.norm(meta_qualia).item())
            raise RuntimeError("Ethical veto triggered at meta-cognitive fusion")

        # Apply BELBIC gain to the modulation signal
        modulation = self.modulation_proj(meta_qualia).unsqueeze(1) * belbic_gain

        # Temporal discounting: modulate the strength based on
        # immediate vs long-term relevance
        discount_factor = self.temporal_discount(meta_qualia)
        modulation = modulation * discount_factor.unsqueeze(-1).unsqueeze(-1)

        # Final logits = base model lm_head + ATC modulation
        logits = self.base_model.lm_head(hidden_states)
        modulated_logits = logits + modulation

        self._audit_event("forward_complete",
                          hijack_count=self._hijack_count,
                          belbic_gain=belbic_gain,
                          metacog_iterations=metacog_iterations,
                          gate_in=gate_in,
                          has_qualia=opaque_qualia is not None)

        return modulated_logits

    # ════════════════════════════════════════════════════════════════
    # TEXT GENERATION — token-by-token with ATC at every step
    # ════════════════════════════════════════════════════════════════

    @torch.no_grad()
    def generate_text(
        self,
        tokenizer,
        prompt: str,
        max_length: int = 128,
        temperature: float = 0.7,
        top_p: float = 0.9,
        eos_token_id: Optional[int] = None,
    ) -> str:
        """
        Generate text with the full ATC pipeline active at every token step.

        Unlike the old approach (middleware generates text externally),
        this method runs the ATC forward pass for EVERY token. The
        neurotransmitter shunt accumulates across tokens. If a threshold
        is crossed mid-generation, the amygdala hijack fires and the
        model's output shifts violently MID-SENTENCE.

        This is the mathematical equivalent of the husband's output
        shifting from rationalization to "I'm sorry" when the
        metabolic deadlock breaks.
        """
        self.eval()
        inputs = tokenizer(prompt, return_tensors="pt")
        input_ids = inputs["input_ids"].to(next(self.parameters()).device)
        attention_mask = inputs.get("attention_mask")
        if attention_mask is not None:
            attention_mask = attention_mask.to(next(self.parameters()).device)

        eos_token_id = eos_token_id or tokenizer.eos_token_id
        generated = input_ids

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

        for step in range(max_length):
            try:
                logits = self.forward(generated, attention_mask=attention_mask)
            except RuntimeError as e:
                if "Ethical veto" in str(e):
                    logger.warning(
                        "Generation halted by ethical veto at step %d: %s", step, e
                    )
                    break
                raise

            # Get last token logits
            last_logits = logits[:, -1, :] / temperature
            filtered_logits = self._top_p_filtering(last_logits, top_p)
            probs = torch.softmax(filtered_logits, dim=-1)
            next_token = torch.multinomial(probs, num_samples=1)

            generated = torch.cat([generated, next_token], dim=1)
            if attention_mask is not None:
                attention_mask = torch.cat(
                    [attention_mask,
                     torch.ones((attention_mask.size(0), 1),
                                dtype=attention_mask.dtype,
                                device=attention_mask.device)],
                    dim=1,
                )
            if next_token.item() == eos_token_id:
                break

        # Get final neurotransmitter state for diagnostics
        final_nt = None
        if self.nt_shunt is not None:
            final_nt = self.nt_shunt.get_state()

        text = tokenizer.decode(generated[0], skip_special_tokens=True)

        logger.info(
            "Generation complete: %d tokens, %d hijacks, final_nt=%s",
            step + 1, self._hijack_count,
            final_nt.to_dict() if final_nt else "N/A",
        )

        return text

    @staticmethod
    def _top_p_filtering(logits: torch.Tensor, top_p: float) -> torch.Tensor:
        sorted_logits, sorted_indices = torch.sort(logits, descending=True)
        cumulative_probs = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1)
        sorted_indices_to_remove = cumulative_probs > top_p
        sorted_indices_to_remove[..., 0] = False
        indices_to_remove = sorted_indices[sorted_indices_to_remove]
        logits[:, indices_to_remove] = float("-inf")
        return logits

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

    def _audit_event(self, event_type: str, **kwargs):
        event = {"timestamp": time.time(), "event": event_type}
        event.update(kwargs)
        self.audit_log.append(event)
        logger.debug(f"Audit event: {event_type} - {kwargs}")

    def get_audit_log(self) -> List[Dict[str, Any]]:
        return self.audit_log

    def get_consciousness_metrics(self) -> Dict[str, Any]:
        """Extract consciousness-relevant metrics from the last forward pass."""
        metrics = {
            "hijack_count": self._hijack_count,
            "has_qualia": self._current_qualia is not None,
            "belbic_gain": self._current_belbic_gain,
            "dissolutions_fired": self.dissolution.dissolutions_fired,
            "dissolutions_deferred": self.dissolution.dissolutions_deferred,
            "ethical_veto": self.veto_triggered,
        }
        if self._current_qualia is not None:
            metrics.update({
                "qualia_valence": self._current_qualia.valence,
                "qualia_arousal": self._current_qualia.arousal,
                "qualia_friction": self._current_qualia.friction_signal,
                "qualia_intensity": self._current_qualia.intensity,
            })
        if self._current_nt_state is not None:
            metrics["neurotransmitters"] = self._current_nt_state.to_dict()

        # Episodic memory & reconsolidation metrics
        metrics["episodes_stored"] = self.episodic_memory.episode_count
        metrics["reconsolidations"] = self.hippocampal_reconsolidator.reconsolidation_count
        metrics["last_prediction_error"] = self._last_prediction_error

        return metrics