tests/test_atc_pipeline.py

"""
Comprehensive automated tests for the NIMA Unified Model's ATC pipeline.

Validates that the ATC cognitive pipeline modulates the model's output
differently from a bare forward pass.  All tests use a mock base model
and never load the real Phi-4-mini weights.
"""

import pytest
import torch
import torch.nn as nn
from unittest.mock import MagicMock, patch

from nima_unified.core.neurotransmitter_shunt import (
    NeurotransmitterShunt,
    NeurotransmitterState,
    IDX_NOREPINEPHRINE,
    IDX_CORTISOL,
    IDX_DOPAMINE,
    IDX_ADENOSINE,
    NOREPINEPHRINE_DECAY,
    CORTISOL_DECAY,
    DOPAMINE_DECAY,
    ADENOSINE_DECAY,
    FRICTION_CORTISOL_INJECT,
    FRICTION_ADENOSINE_INJECT,
    NOVELTY_NE_INJECT,
    METACOGNITIVE_LOOP_CORTISOL_INJECT,
    METACOGNITIVE_LOOP_ADENOSINE_INJECT,
)
from nima_unified.core.deep_surgery import (
    OpaqueQualiaSignature,
    TRNPredictiveGate,
    DissolutionModule,
    BELBICDualPathway,
    MetacognitiveLoopModule,
    IrrationalSparkModule,
    EpisodicMemoryModule,
    HippocampalReconsolidator,
    EthicalGuardian,
    ATCDeepSurgery,
)


# ── Helper: lightweight mock transformer layer ────────────────────────

class MockTransformerLayer(nn.Module):
    """Minimal transformer-style layer: Linear + residual, fast enough for
    tests but with real gradient plumbing."""

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

    def forward(self, hidden_states, attention_mask=None, **kwargs):
        return (hidden_states + self.linear(hidden_states),)


def make_mock_base_model(hidden_size=256, num_layers=12, vocab_size=1000):
    """Build a mock model that satisfies the ATCDeepSurgery interface."""
    model = MagicMock()
    model.config = MagicMock()
    model.config.hidden_size = hidden_size
    model.config.vocab_size = vocab_size

    layers = nn.ModuleList(
        [MockTransformerLayer(hidden_size) for _ in range(num_layers)]
    )
    model.model = MagicMock()
    model.model.layers = layers
    model.get_input_embeddings.return_value = nn.Embedding(vocab_size, hidden_size)
    model.lm_head = nn.Linear(hidden_size, vocab_size, bias=False)

    return model, layers


# ══════════════════════════════════════════════════════════════════════
# TestNeurotransmitterShunt  (~15 tests)
# ══════════════════════════════════════════════════════════════════════

class TestNeurotransmitterShunt:

    def test_initial_state_is_zero(self):
        shunt = NeurotransmitterShunt()
        state = shunt.get_state()
        assert state.norepinephrine == 0.0
        assert state.cortisol == 0.0
        assert state.dopamine == 0.0
        assert state.adenosine == 0.0

    def test_inject_friction_raises_cortisol_and_adenosine(self):
        shunt = NeurotransmitterShunt()
        shunt.inject_friction(1.0)
        state = shunt.get_state()
        assert state.cortisol == pytest.approx(FRICTION_CORTISOL_INJECT)
        assert state.adenosine == pytest.approx(FRICTION_ADENOSINE_INJECT)

    def test_inject_norepinephrine_caps_at_one(self):
        shunt = NeurotransmitterShunt()
        for _ in range(100):
            shunt.inject_norepinephrine(NOVELTY_NE_INJECT)
        state = shunt.get_state()
        assert state.norepinephrine <= 1.0
        assert state.norepinephrine == pytest.approx(1.0)

    def test_inject_dopamine(self):
        shunt = NeurotransmitterShunt()
        shunt.inject_dopamine(0.5)
        state = shunt.get_state()
        assert state.dopamine == pytest.approx(0.5)

    def test_decay_reduces_values(self):
        shunt = NeurotransmitterShunt()
        shunt.inject_norepinephrine(0.5)
        shunt.inject_adenosine_via_friction = lambda: None  # no-op
        # Manually set adenosine so we can compare decay rates
        shunt._vector[IDX_ADENOSINE] = 0.5
        state_before = shunt.read_and_decay(dt=1.0)
        state_after = shunt.read_and_decay(dt=1.0)
        # Both should have decreased
        assert state_after.norepinephrine < state_before.norepinephrine
        assert state_after.adenosine < state_before.adenosine
        # NE should decay faster (higher decay rate)
        ne_ratio = state_after.norepinephrine / max(state_before.norepinephrine, 1e-9)
        ad_ratio = state_after.adenosine / max(state_before.adenosine, 1e-9)
        assert ne_ratio < ad_ratio

    def test_decay_order_ne_fastest_adenosine_slowest(self):
        """Verify half-lives: NE: 3.0, Cortisol: 0.1, Dopamine: 0.8, Adenosine: 0.05"""
        # Exponential decay: after 1 second, value *= exp(-rate * dt)
        # Higher rate = faster decay
        assert NOREPINEPHRINE_DECAY == 3.0
        assert CORTISOL_DECAY == 0.1
        assert DOPAMINE_DECAY == 0.8
        assert ADENOSINE_DECAY == 0.05
        # Ordering: NE (3.0) > Dopamine (0.8) > Cortisol (0.1) > Adenosine (0.05)
        assert NOREPINEPHRINE_DECAY > DOPAMINE_DECAY
        assert DOPAMINE_DECAY > CORTISOL_DECAY
        assert CORTISOL_DECAY > ADENOSINE_DECAY

    def test_adenosine_floor_from_metabolic_reserve(self):
        shunt = NeurotransmitterShunt()
        shunt.update_metabolic_reserve(0.2)
        # adenosine_floor = (1.0 - metabolic_reserve) * 0.5 = 0.8 * 0.5 = 0.4
        state = shunt.read_and_decay(dt=0.001)
        assert state.adenosine == pytest.approx(0.4, abs=0.01)

    def test_auto_hijack_at_threshold(self):
        shunt = NeurotransmitterShunt()
        # Bypass the inject cap and set adenosine directly past 0.95
        shunt._vector[IDX_ADENOSINE] = 0.96
        state = shunt.read_and_decay(dt=0.001)
        assert state.hijack_active is True
        assert "ADENOSINE" in state.hijack_reason

    def test_hijack_clears_after_timeout(self):
        shunt = NeurotransmitterShunt()
        shunt._vector[IDX_ADENOSINE] = 0.96
        state = shunt.read_and_decay(dt=0.001)
        assert state.hijack_active is True
        # Patch time so hijack appears to have fired > 0.1s ago
        with patch("nima_unified.core.neurotransmitter_shunt.time.time") as mock_time:
            mock_time.return_value = shunt._hijack_timestamp + 0.2
            state2 = shunt.read_and_decay(dt=0.001)
            assert state2.hijack_active is False

    def test_suppression_trigger(self):
        shunt = NeurotransmitterShunt()
        result = shunt.suppress_metabolic_and_trigger_hijack("test_suppress")
        assert result is True
        state = shunt.get_state()
        assert state.hijack_active is True
        assert state.hijack_reason == "test_suppress"
        assert state.norepinephrine > 0.0

    def test_reset_clears_everything(self):
        shunt = NeurotransmitterShunt()
        shunt.inject_friction(1.0)
        shunt.inject_norepinephrine(0.5)
        shunt.inject_dopamine(0.5)
        shunt.suppress_metabolic_and_trigger_hijack("reset_test")
        shunt.reset()
        state = shunt.get_state()
        assert state.norepinephrine == 0.0
        assert state.cortisol == 0.0
        assert state.dopamine == 0.0
        assert state.adenosine == 0.0
        assert state.hijack_active is False

    def test_power_spike_prediction_raises_cortisol(self):
        shunt = NeurotransmitterShunt()
        shunt.set_power_spike_predicted(True)
        state = shunt.get_state()
        assert state.power_spike_predicted is True
        assert state.cortisol == pytest.approx(
            FRICTION_CORTISOL_INJECT * 0.5
        )

    def test_get_state_returns_snapshot(self):
        shunt = NeurotransmitterShunt()
        shunt.inject_friction(1.0)
        shunt.inject_norepinephrine(0.3)
        shunt.inject_dopamine(0.2)
        state = shunt.get_state()
        assert isinstance(state, NeurotransmitterState)
        assert state.norepinephrine == 0.3
        assert state.cortisol == FRICTION_CORTISOL_INJECT
        assert state.dopamine == 0.2
        assert state.total_firings > 0
        assert state.last_update_ts > 0.0
        assert isinstance(state.hijack_active, bool)
        assert isinstance(state.hijack_reason, str)
        assert isinstance(state.metabolic_reserve, float)
        assert isinstance(state.power_spike_predicted, bool)

    def test_inject_metacognitive_strain(self):
        shunt = NeurotransmitterShunt()
        shunt.inject_metacognitive_strain(loop_iterations=3, loop_stress=0.5)
        state = shunt.get_state()
        assert state.cortisol == pytest.approx(
            METACOGNITIVE_LOOP_CORTISOL_INJECT * 0.5, abs=1e-6
        )
        assert state.adenosine == pytest.approx(
            METACOGNITIVE_LOOP_ADENOSINE_INJECT * 3, abs=1e-6
        )

    def test_history_recording(self):
        shunt = NeurotransmitterShunt()
        assert len(shunt._history) == 0
        shunt.inject_friction(1.0)
        # inject methods do NOT record history directly; only read_and_decay does
        assert len(shunt._history) == 0
        # read_and_decay records a "read" event
        shunt.read_and_decay(dt=0.001)
        assert len(shunt._history) == 1
        assert shunt._history[-1]["event"] == "read"
        # suppress_metabolic_and_trigger_hijack records a "suppression_hijack" event
        shunt.suppress_metabolic_and_trigger_hijack("test")
        assert len(shunt._history) == 2
        assert shunt._history[-1]["event"] == "suppression_hijack"
        # Each entry should have a timestamp
        assert "ts" in shunt._history[0]


# ══════════════════════════════════════════════════════════════════════
# TestTRNPredictiveGate  (~5 tests)
# ══════════════════════════════════════════════════════════════════════

class TestTRNPredictiveGate:

    @pytest.fixture
    def gate(self):
        return TRNPredictiveGate(hidden_size=64)

    @pytest.fixture
    def hidden(self):
        return torch.randn(2, 8, 64)

    def test_high_confidence_gates_out(self, gate, hidden):
        gate_in, confidence = gate(hidden, prediction_confidence=0.99)
        assert gate_in is False

    def test_low_confidence_gates_in(self, gate, hidden):
        # External confidence of 0.0 blended with model confidence -> low
        gate_in, confidence = gate(hidden, prediction_confidence=0.0)
        assert gate_in is True

    def test_output_shape(self, gate, hidden):
        result = gate(hidden, prediction_confidence=0.5)
        assert isinstance(result, tuple)
        assert len(result) == 2
        assert isinstance(result[0], bool)
        assert isinstance(result[1], float)

    def test_learnable_threshold(self, gate):
        assert isinstance(gate.gate_threshold, nn.Parameter)
        assert gate.gate_threshold.requires_grad

    def test_confidence_blending(self, gate, hidden):
        """Model confidence and external confidence are blended 50/50."""
        # With external confidence = 1.0, blended = 0.5 * model_conf + 0.5 * 1.0
        # Even if model confidence is very low, the blend is at least 0.5
        # which is above the default threshold (0.15) -> gate_in=False
        gate_in, confidence = gate(hidden, prediction_confidence=1.0)
        assert gate_in is False
        assert confidence > 0.4  # 0.5 * model + 0.5 * 1.0 > 0.4


# ══════════════════════════════════════════════════════════════════════
# TestDissolutionModule  (~6 tests)
# ══════════════════════════════════════════════════════════════════════

class TestDissolutionModule:

    @pytest.fixture
    def dissolution(self):
        return DissolutionModule(hidden_size=64, qualia_dim=32)

    @pytest.fixture
    def hidden(self):
        return torch.randn(2, 8, 64)

    def test_gate_out_returns_none_qualia(self, dissolution, hidden):
        qualia, offset, fired = dissolution(hidden, gate_in=False)
        assert qualia is None
        assert fired is False
        assert offset.shape == hidden.shape
        assert offset.abs().max() == 0.0

    def test_gate_in_returns_qualia_signature(self, dissolution, hidden):
        # Force alpha phase to always allow dissolution
        dissolution.check_alpha_phase = MagicMock(return_value=True)
        qualia, offset, fired = dissolution(hidden, gate_in=True)
        assert qualia is not None
        assert isinstance(qualia, OpaqueQualiaSignature)
        assert fired is True

    def test_dissolution_offset_nonzero_when_fired(self, dissolution, hidden):
        dissolution.check_alpha_phase = MagicMock(return_value=True)
        _, offset, fired = dissolution(hidden, gate_in=True)
        assert fired is True
        assert offset.abs().sum() > 0.0

    def test_qualia_values_bounded(self, dissolution, hidden):
        dissolution.check_alpha_phase = MagicMock(return_value=True)
        qualia, _, _ = dissolution(hidden, gate_in=True)
        # valence is raw from Tanh: [-1, 1]
        assert -1.0 <= qualia.valence <= 1.0
        # arousal, intensity, friction_signal, memory_salience are mapped to [0, 1]
        assert 0.0 <= qualia.arousal <= 1.0
        assert 0.0 <= qualia.intensity <= 1.0
        assert 0.0 <= qualia.friction_signal <= 1.0
        assert 0.0 <= qualia.memory_salience <= 1.0

    def test_alpha_phase_deferral(self, dissolution, hidden):
        dissolution.check_alpha_phase = MagicMock(return_value=False)
        qualia, offset, fired = dissolution(hidden, gate_in=True)
        assert qualia is None
        assert fired is False
        assert dissolution.dissolutions_deferred > 0

    def test_dissolution_fired_counter(self, dissolution, hidden):
        dissolution.check_alpha_phase = MagicMock(return_value=True)
        assert dissolution.dissolutions_fired == 0
        dissolution(hidden, gate_in=True)
        assert dissolution.dissolutions_fired == 1
        dissolution(hidden, gate_in=True)
        assert dissolution.dissolutions_fired == 2


# ══════════════════════════════════════════════════════════════════════
# TestBELBICDualPathway  (~5 tests)
# ══════════════════════════════════════════════════════════════════════

class TestBELBICDualPathway:

    @pytest.fixture
    def belbic(self):
        return BELBICDualPathway(hidden_size=64)

    @pytest.fixture
    def sensory(self):
        return torch.tensor([[0.5, 0.5, 0.5, 0.5]])

    def test_initial_zero_output(self, belbic, sensory):
        amygdala_out, ofc_out, gain = belbic(sensory)
        # Weights initialized to zero → sigmoid(0) = 0.5
        assert amygdala_out == pytest.approx(0.5)
        assert ofc_out == pytest.approx(0.5)
        assert gain == pytest.approx(1.0, abs=0.05)

    def test_gain_within_bounds(self, belbic):
        """gain is always in [GAIN_FLOOR, GAIN_CEIL] = [0.2, 2.0]."""
        for _ in range(20):
            s = torch.rand(1, 4)
            _, _, gain = belbic(s)
            assert BELBICDualPathway.GAIN_FLOOR <= gain <= BELBICDualPathway.GAIN_CEIL

    def test_amygdala_strengthened_on_positive_reward(self, belbic, sensory):
        old_weights = belbic.amygdala.weight.data.clone()
        belbic.update(sensory, reward=1.0)
        # Amygdala weights should have changed (increased in magnitude)
        diff = (belbic.amygdala.weight.data - old_weights).abs().sum().item()
        assert diff > 0.0

    def test_ofc_strengthened_on_negative_reward(self, belbic, sensory):
        old_weights = belbic.ofc.weight.data.clone()
        belbic.update(sensory, reward=-1.0)
        # OFC weights should have changed (strengthened inhibition)
        diff = (belbic.ofc.weight.data - old_weights).abs().sum().item()
        assert diff > 0.0

    def test_output_tuple_shape(self, belbic, sensory):
        result = belbic(sensory)
        assert isinstance(result, tuple)
        assert len(result) == 3
        assert all(isinstance(v, float) for v in result)


# ══════════════════════════════════════════════════════════════════════
# TestMetacognitiveLoopModule  (~5 tests)
# ══════════════════════════════════════════════════════════════════════

class TestMetacognitiveLoopModule:

    @pytest.fixture
    def metacog(self):
        return MetacognitiveLoopModule(hidden_size=64)

    @pytest.fixture
    def hidden(self):
        return torch.randn(2, 8, 64)

    def test_high_comprehension_no_loop(self, metacog, hidden):
        """When comprehension > 0.7, iterations = 0."""
        # Patch comprehension_head to always return high value
        metacog.comprehension_head = nn.Sequential(
            nn.Linear(64, 128), nn.ReLU(), nn.Linear(128, 1), nn.Sigmoid()
        )
        # Force high comprehension by making the last layer output large positive
        with torch.no_grad():
            metacog.comprehension_head[-2].weight.fill_(0.0)
            metacog.comprehension_head[-2].bias.fill_(10.0)  # sigmoid(10) ≈ 1.0
        _, iters, stress, strain, spark = metacog(hidden)
        assert iters == 0

    def test_low_comprehension_enters_loop(self, metacog, hidden):
        """When comprehension < 0.7, iterations > 0."""
        # Force low comprehension
        with torch.no_grad():
            metacog.comprehension_head[-2].weight.fill_(0.0)
            metacog.comprehension_head[-2].bias.fill_(-10.0)  # sigmoid(-10) ≈ 0.0
        _, iters, stress, strain, spark = metacog(hidden)
        assert iters > 0

    def test_deadlock_triggers_spark(self, metacog, hidden):
        """When stress > 0.6 and iterations > 3, spark_fired = True."""
        # Force low comprehension so loop runs, and high strain
        with torch.no_grad():
            metacog.comprehension_head[-2].weight.fill_(0.0)
            metacog.comprehension_head[-2].bias.fill_(-10.0)  # low comprehension
            metacog.strain_head[-2].weight.fill_(0.0)
            metacog.strain_head[-2].bias.fill_(10.0)  # sigmoid(10) ≈ 1.0 (high strain)
        _, iters, stress, strain, spark = metacog(hidden)
        # With strain=1.0, stress = strain * (iters / MAX_ITERATIONS)
        # At iteration 4: stress = 1.0 * (4/5) = 0.8 > 0.6 and iters > 3
        if iters > 3:
            assert spark is True
        else:
            # If the loop breaks early (shouldn't with low comprehension), still ok
            assert iters >= 0

    def test_strain_always_returned(self, metacog, hidden):
        """strain is always a float in [0, 1]."""
        _, _, _, strain, _ = metacog(hidden)
        assert isinstance(strain, float)
        assert 0.0 <= strain <= 1.0

    def test_max_iterations_limit(self, metacog, hidden):
        """Never exceeds MAX_ITERATIONS = 5."""
        with torch.no_grad():
            metacog.comprehension_head[-2].weight.fill_(0.0)
            metacog.comprehension_head[-2].bias.fill_(-10.0)
        _, iters, _, _, _ = metacog(hidden)
        assert 0 <= iters <= MetacognitiveLoopModule.MAX_ITERATIONS


# ══════════════════════════════════════════════════════════════════════
# TestIrrationalSparkModule  (~5 tests)
# ══════════════════════════════════════════════════════════════════════

class TestIrrationalSparkModule:

    @pytest.fixture
    def spark(self):
        return IrrationalSparkModule(hidden_size=64, vocab_size=100)

    @pytest.fixture
    def hidden(self):
        torch.manual_seed(42)
        return torch.randn(2, 8, 64)

    @pytest.fixture
    def nominal_nt_state(self):
        return NeurotransmitterState(
            norepinephrine=0.1,
            cortisol=0.2,
            dopamine=0.1,
            adenosine=0.1,
            hijack_active=False,
        )

    def test_no_hijack_when_nominal(self, spark, hidden, nominal_nt_state):
        modulated, hijack_fired, reason = spark(hidden, None, nominal_nt_state)
        assert hijack_fired is False
        assert reason == ""
        assert torch.equal(modulated, hidden)

    def test_hijack_on_adenosine_critical(self, spark, hidden):
        nt = NeurotransmitterState(adenosine=0.96, hijack_active=False)
        modulated, hijack_fired, reason = spark(hidden, None, nt)
        assert hijack_fired is True
        assert "ADENOSINE" in reason

    def test_hijack_on_cortisol_critical(self, spark, hidden):
        nt = NeurotransmitterState(cortisol=0.97, hijack_active=False)
        modulated, hijack_fired, reason = spark(hidden, None, nt)
        assert hijack_fired is True
        assert "CORTISOL" in reason

    def test_hijack_on_qualia_crisis(self, spark, hidden, nominal_nt_state):
        qualia = OpaqueQualiaSignature(
            valence=0.0, arousal=0.9, intensity=0.8,
            friction_signal=0.9, memory_salience=0.5,
        )
        modulated, hijack_fired, reason = spark(hidden, qualia, nominal_nt_state)
        assert hijack_fired is True
        assert "QUALIA_CRISE" in reason

    def test_hidden_states_modified_on_hijack(self, spark, hidden):
        nt = NeurotransmitterState(adenosine=0.96, hijack_active=False)
        modulated, hijack_fired, _ = spark(hidden, None, nt)
        assert hijack_fired is True
        assert not torch.equal(modulated, hidden)
        # The difference should be nonzero
        assert (modulated - hidden).abs().sum() > 0.0


# ══════════════════════════════════════════════════════════════════════
# TestEpisodicMemoryModule  (~5 tests)
# ══════════════════════════════════════════════════════════════════════

class TestEpisodicMemoryModule:

    @pytest.fixture
    def memory(self):
        return EpisodicMemoryModule(hidden_size=64, max_episodes=10, embedding_dim=32)

    @pytest.fixture
    def hidden(self):
        torch.manual_seed(99)
        return torch.randn(1, 4, 64)

    def test_empty_memory_returns_zero_prediction_error(self, memory, hidden):
        _, pe, best_match = memory(hidden)
        assert pe == 0.0
        assert best_match is None

    def test_store_and_retrieve(self, memory, hidden):
        # Store an episode
        qualia = OpaqueQualiaSignature(valence=0.5, arousal=0.5, intensity=0.5,
                                       friction_signal=0.3, memory_salience=0.4)
        memory.store_episode(hidden, qualia_signature=qualia)
        assert memory.episode_count == 1

        # Query with similar hidden states → should match
        retrieval_signal, pe, best_match = memory(hidden)
        assert best_match is not None
        assert pe < 1.0
        assert "similarity" in best_match

    def test_prediction_error_high_for_different_inputs(self, memory, hidden):
        # Store with one hidden state
        memory.store_episode(hidden)
        # Query with very different hidden state
        different_hidden = torch.randn(1, 4, 64) * 10.0
        _, pe, _ = memory(different_hidden)
        # Very different input should yield high prediction error
        assert pe > 0.5

    def test_episode_count_increments(self, memory, hidden):
        assert memory.episode_count == 0
        memory.store_episode(hidden)
        assert memory.episode_count == 1
        memory.store_episode(hidden)
        assert memory.episode_count == 2
        memory.store_episode(hidden)
        assert memory.episode_count == 3

    def test_wraps_at_max_episodes(self, memory, hidden):
        max_ep = memory.max_episodes  # 10
        for i in range(max_ep + 5):
            memory.store_episode(hidden)
        assert memory.episode_count == max_ep + 5
        # Verify the embedding at slot 0 was overwritten
        slot = 0
        # Store once more to fill slot 0, then check
        memory.store_episode(hidden)
        # Count should keep incrementing; slot wraps via modulo
        assert memory.episode_count == max_ep + 6
        # Slot for the next store should wrap
        assert (memory.episode_count) % max_ep != memory.episode_count


# ══════════════════════════════════════════════════════════════════════
# TestHippocampalReconsolidator  (~4 tests)
# ══════════════════════════════════════════════════════════════════════

class TestHippocampalReconsolidator:

    @pytest.fixture
    def recon(self):
        return HippocampalReconsolidator()

    @pytest.fixture
    def hidden(self):
        return torch.randn(1, 4, 64)

    def test_no_reconsolidation_below_threshold(self, recon, hidden):
        result = recon(
            hidden,
            episode_valence=0.5,
            episode_arousal=0.3,
            prediction_error=0.2,
            current_valence=-0.5,
            current_arousal=0.9,
        )
        reconsolidated, new_v, new_a, reason = result
        assert reconsolidated is False
        assert "below" in reason

    def test_reconsolidation_above_threshold(self, recon, hidden):
        result = recon(
            hidden,
            episode_valence=0.5,
            episode_arousal=0.3,
            prediction_error=0.8,
            current_valence=-0.5,
            current_arousal=0.9,
        )
        reconsolidated, new_v, new_a, reason = result
        assert reconsolidated is True
        assert isinstance(new_v, float)
        assert isinstance(new_a, float)
        assert isinstance(reason, str)
        assert "reconsolidated" in reason

    def test_valence_shifts_toward_current(self, recon, hidden):
        """Reconsolidated valence is a blend of old and new, clamped to [-1, 1]."""
        torch.manual_seed(7)
        episode_valence = 0.8
        current_valence = -0.8
        pred_error = 0.9

        _, new_v, _, _ = recon(
            hidden,
            episode_valence=episode_valence,
            episode_arousal=0.3,
            prediction_error=pred_error,
            current_valence=current_valence,
            current_arousal=0.5,
        )
        # new = old * 0.7 + projection * 0.3 + noise, clamped to [-1, 1]
        assert isinstance(new_v, float)
        assert -1.0 <= new_v <= 1.0
        # The 70/30 blend with projection and noise means the value
        # should differ from the pure episode_valence (noise alone
        # almost guarantees this, and the projection is nonzero)
        assert new_v != pytest.approx(episode_valence, abs=0.05)

    def test_reconsolidation_count_increments(self, recon, hidden):
        assert recon.reconsolidation_count == 0
        recon(
            hidden,
            episode_valence=0.5,
            episode_arousal=0.3,
            prediction_error=0.8,
            current_valence=-0.5,
            current_arousal=0.9,
        )
        assert recon.reconsolidation_count == 1
        recon(
            hidden,
            episode_valence=0.1,
            episode_arousal=0.2,
            prediction_error=0.9,
            current_valence=0.8,
            current_arousal=0.7,
        )
        assert recon.reconsolidation_count == 2


# ══════════════════════════════════════════════════════════════════════
# TestATCDeepSurgeryIntegration  (~6 tests)
# ══════════════════════════════════════════════════════════════════════

class TestATCDeepSurgeryIntegration:

    @pytest.fixture
    def setup(self):
        """Create a fully wired ATC deep surgery with mock base model."""
        torch.manual_seed(123)
        hidden_size = 256
        num_layers = 12
        vocab_size = 1000

        mock_model, real_layers = make_mock_base_model(
            hidden_size=hidden_size, num_layers=num_layers, vocab_size=vocab_size
        )

        # Replace the mock layers list with the real nn.ModuleList
        mock_model.model.layers = real_layers
        mock_model.get_input_embeddings.return_value = nn.Embedding(vocab_size, hidden_size)
        mock_model.lm_head = nn.Linear(hidden_size, vocab_size, bias=False)

        shunt = NeurotransmitterShunt()
        guardian = EthicalGuardian(threshold=100.0)  # High threshold to avoid vetoes

        surgery = ATCDeepSurgery(
            base_model=mock_model,
            ethical_guardian=guardian,
            num_layers=num_layers,
            qualia_dim=64,
            neurotransmitter_shunt=shunt,
        )

        return surgery, shunt, hidden_size, vocab_size

    def test_forward_returns_modulated_logits(self, setup):
        surgery, shunt, hidden_size, vocab_size = setup
        batch_size, seq_len = 2, 8
        input_ids = torch.randint(0, vocab_size, (batch_size, seq_len))

        logits = surgery(input_ids)

        assert logits.shape == (batch_size, seq_len, vocab_size)
        assert logits.dtype == torch.float32

    def test_atc_modulation_differs_from_bare_model(self, setup):
        """ATC modulated logits should differ from bare model logits."""
        surgery, shunt, hidden_size, vocab_size = setup
        batch_size, seq_len = 2, 8
        input_ids = torch.randint(0, vocab_size, (batch_size, seq_len))

        # ── ATC forward ──
        with torch.no_grad():
            atc_logits = surgery(input_ids)

        # ── Bare model forward (same layers + lm_head, no ATC) ──
        embeddings = surgery.base_model.get_input_embeddings()(input_ids)
        h = embeddings
        for layer in surgery._get_transformer_layers():
            h = layer(h)[0]
        bare_logits = surgery.base_model.lm_head(h)

        # They should NOT be identical
        assert not torch.allclose(atc_logits, bare_logits, atol=0.0)

    def test_neurotransmitter_shunt_receives_signals(self, setup):
        """After forward, shunt.total_firings > 0."""
        surgery, shunt, hidden_size, vocab_size = setup
        input_ids = torch.randint(0, vocab_size, (1, 8))

        with torch.no_grad():
            surgery(input_ids)

        state = shunt.get_state()
        assert state.total_firings > 0

    def test_consciousness_metrics_populated(self, setup):
        """get_consciousness_metrics returns expected keys."""
        surgery, shunt, hidden_size, vocab_size = setup
        input_ids = torch.randint(0, vocab_size, (1, 8))

        with torch.no_grad():
            surgery(input_ids)

        metrics = surgery.get_consciousness_metrics()
        expected_keys = {
            "hijack_count", "has_qualia", "belbic_gain",
            "dissolutions_fired", "dissolutions_deferred",
            "ethical_veto", "episodes_stored",
            "reconsolidations", "last_prediction_error",
        }
        for key in expected_keys:
            assert key in metrics, f"Missing key: {key}"

    def test_ethical_veto_raises_runtime_error(self, setup):
        """A qualia vector with very high norm triggers RuntimeError."""
        # Build a surgery with a very low ethical threshold
        torch.manual_seed(456)
        hidden_size = 256
        num_layers = 12
        vocab_size = 1000
        mock_model, real_layers = make_mock_base_model(
            hidden_size=hidden_size, num_layers=num_layers, vocab_size=vocab_size
        )
        mock_model.model.layers = real_layers
        mock_model.get_input_embeddings.return_value = nn.Embedding(vocab_size, hidden_size)
        mock_model.lm_head = nn.Linear(hidden_size, vocab_size, bias=False)

        shunt = NeurotransmitterShunt()
        # Very low threshold — the tanh-bounded qualia vector norm is at most ~sqrt(5)
        # (5 dims, each in [-1,1]), so threshold=0.1 should trigger veto if
        # the dissolution output is nonzero.
        # Actually, qualia.to_tensor gives [valence, arousal, intensity, friction, memory_salience]
        # where arousal/intensity/friction/memory_salience are in [0,1]. The norm of [0,0,0,0,0]
        # is 0. We need the dissolution to fire and produce a nonzero qualia.
        # The dissolution encoder produces values in [-1,1] for valence and
        # [0,1] for the mapped fields. The qualia tensor norm is bounded.
        # Use threshold=0.0 to guarantee veto.
        guardian = EthicalGuardian(threshold=0.0)

        surgery = ATCDeepSurgery(
            base_model=mock_model,
            ethical_guardian=guardian,
            num_layers=num_layers,
            qualia_dim=64,
            neurotransmitter_shunt=shunt,
        )

        input_ids = torch.randint(0, vocab_size, (1, 8))
        with pytest.raises(RuntimeError, match="Ethical veto"):
            with torch.no_grad():
                surgery(input_ids)

    def test_hijack_count_tracking(self, setup):
        """When conditions trigger hijack, _hijack_count > 0."""
        surgery, shunt, hidden_size, vocab_size = setup
        input_ids = torch.randint(0, vocab_size, (1, 8))

        # Manually push adenosine past the critical threshold so the
        # irrational spark fires in Layer 5
        shunt._vector[IDX_ADENOSINE] = 0.96

        with torch.no_grad():
            surgery(input_ids)

        assert surgery._hijack_count > 0