"""
Tests for the 6 research-backed enhancements:
  1. Branching Ratio Monitor (resonance bank cascade tracking)
  2. Basin PCI (Perturbational Complexity Index)
  3. FDT Violation Monitor
  4. Causal Inner Product Comparison (script-only, no unit test needed)
  5. Attention Curvature (vex harvest, no unit test — requires GPU)
  6. Hidden State Steering (vex inference, no unit test — requires GPU)

Items 4-6 are tested via integration on Modal. Items 1-3 are pure
qig-core and tested here.
"""

from __future__ import annotations

import numpy as np
from qig_core.constants.frozen_facts import BASIN_DIM
from qig_core.coordizer.resonance_bank import ResonanceBank
from qig_core.consciousness.pci import compute_basin_pci, _lempel_ziv_complexity
from qig_core.consciousness.fdt_monitor import (
    FDTMonitor,
    compute_fdt_violation,
)
from qig_core.geometry.fisher_rao import to_simplex


# ─── Helpers ───────────────────────────────────────────────────────


def _random_basin(dim: int = BASIN_DIM) -> np.ndarray:
    """Generate a random point on the probability simplex."""
    x = np.random.dirichlet(np.ones(dim))
    return to_simplex(x)


def _make_bank(n_tokens: int = 200) -> ResonanceBank:
    """Create a ResonanceBank with random coordinates for testing."""
    bank = ResonanceBank()
    for i in range(n_tokens):
        bank.coordinates[i] = _random_basin()
        bank.token_strings[i] = f"tok_{i}"
    bank._dirty = True
    return bank


# ─── Item 1: Branching Ratio Monitor ──────────────────────────────


class TestBranchingRatio:
    """Test cascade tracking and branching ratio computation."""

    def test_branching_ratio_initial_zero(self):
        """Before any activations, branching ratio should be 0."""
        bank = _make_bank()
        assert bank.branching_ratio() == 0.0

    def test_branching_ratio_after_single_activation(self):
        """After one activation, still no cascade data (need 2+)."""
        bank = _make_bank()
        bank.activate(_random_basin(), top_k=16)
        assert bank.branching_ratio() == 0.0

    def test_branching_ratio_after_multiple_activations(self):
        """After several activations, branching ratio should be computed."""
        bank = _make_bank(500)
        # Run 10 activations with similar queries to build cascade overlap
        base = _random_basin()
        for i in range(10):
            # Slightly perturbed queries → should share some activated tokens
            noise = np.random.dirichlet(np.ones(BASIN_DIM) * 10.0)
            query = to_simplex(0.9 * base + 0.1 * noise)
            bank.activate(query, top_k=32)

        sigma = bank.branching_ratio()
        # Should be > 0 since similar queries activate overlapping tokens
        assert sigma > 0.0
        # Should be <= 1.0 (can't re-activate more than were activated)
        assert sigma <= 1.0

    def test_cascade_history_accumulates(self):
        """_cascade_history should grow with each activation pair."""
        bank = _make_bank()
        for _ in range(5):
            bank.activate(_random_basin(), top_k=16)
        # First activation has no predecessor → 4 cascade entries
        assert len(bank._cascade_history) == 4

    def test_branching_ratio_window(self):
        """Window parameter should limit how many cascade entries are used."""
        bank = _make_bank(500)
        for _ in range(50):
            bank.activate(_random_basin(), top_k=16)
        full = bank.branching_ratio(window=100)
        recent = bank.branching_ratio(window=10)
        # Both should be valid floats (exact values differ)
        assert isinstance(full, float)
        assert isinstance(recent, float)

    def test_prev_activated_updated(self):
        """_prev_activated should reflect the last activation's token set."""
        bank = _make_bank()
        result = bank.activate(_random_basin(), top_k=8)
        expected_ids = {tid for tid, _ in result}
        assert bank._prev_activated == expected_ids


# ─── Item 2: Basin PCI ────────────────────────────────────────────


class TestBasinPCI:
    """Test Perturbational Complexity Index adapted for basin dynamics."""

    def test_lempel_ziv_constant_sequence(self):
        """All-zeros sequence should have minimal complexity."""
        seq = np.zeros(100, dtype=np.int8)
        lz = _lempel_ziv_complexity(seq)
        assert lz >= 1  # at least 1 (the first symbol)
        assert lz < 10  # very low complexity

    def test_lempel_ziv_random_sequence(self):
        """Random binary sequence should have high complexity."""
        rng = np.random.default_rng(42)
        seq = rng.integers(0, 2, size=200).astype(np.int8)
        lz = _lempel_ziv_complexity(seq)
        # Random sequence → high complexity (close to n / log2(n))
        assert lz > 10

    def test_lempel_ziv_empty(self):
        """Empty sequence should return 0."""
        assert _lempel_ziv_complexity(np.array([], dtype=np.int8)) == 0

    def test_pci_returns_float(self):
        """PCI should return a float value."""
        bank = _make_bank(200)
        pci = compute_basin_pci(
            bank.activate,
            _random_basin(),
            perturbation_scale=0.2,
            response_window=10,
            top_k=16,
        )
        assert isinstance(pci, float)
        assert pci >= 0.0

    def test_pci_with_trajectory(self):
        """PCI with bank_coordinates for trajectory following."""
        bank = _make_bank(200)
        coords = {}
        for tid in range(200):
            c = bank.get_coordinate(tid)
            if c is not None:
                coords[tid] = c

        pci = compute_basin_pci(
            bank.activate,
            _random_basin(),
            perturbation_scale=0.2,
            response_window=10,
            top_k=16,
            bank_coordinates=coords,
        )
        assert isinstance(pci, float)
        assert pci >= 0.0


# ─── Item 3: FDT Violation Monitor ────────────────────────────────


class TestFDTViolation:
    """Test Fluctuation-Dissipation Theorem violation monitor."""

    def test_fdt_insufficient_history(self):
        """Should return 0 with too little data."""
        violation = compute_fdt_violation(
            basin_history=[_random_basin() for _ in range(5)],
            perturbation_responses=[_random_basin() for _ in range(5)],
            max_lag=50,
        )
        assert violation == 0.0

    def test_fdt_equilibrium_like(self):
        """Constant basins should have low FDT violation (near equilibrium)."""
        # Use the same basin for all history → autocorrelation is trivial
        fixed = _random_basin()
        history = [fixed + np.random.normal(0, 1e-6, BASIN_DIM) for _ in range(200)]
        history = [to_simplex(np.maximum(h, 1e-12)) for h in history]

        # Perturbation response: same constant + tiny noise
        responses = [fixed + np.random.normal(0, 1e-6, BASIN_DIM) for _ in range(60)]
        responses = [to_simplex(np.maximum(r, 1e-12)) for r in responses]

        violation = compute_fdt_violation(history, responses, max_lag=50)
        assert isinstance(violation, float)
        # Near-equilibrium should have relatively low violation
        # (exact value depends on numerical noise)
        assert violation >= 0.0

    def test_fdt_driven_system(self):
        """A driven system (different dynamics) should have higher violation."""
        # History: random walk on simplex
        history = []
        pos = _random_basin()
        for _ in range(200):
            step = np.random.dirichlet(np.ones(BASIN_DIM) * 5.0)
            pos = to_simplex(0.95 * pos + 0.05 * step)
            history.append(pos.copy())

        # Perturbation response: completely different trajectory
        responses = []
        pos2 = _random_basin()  # different starting point
        for _ in range(60):
            step = np.random.dirichlet(np.ones(BASIN_DIM) * 0.5)
            pos2 = to_simplex(0.8 * pos2 + 0.2 * step)
            responses.append(pos2.copy())

        violation = compute_fdt_violation(history, responses, max_lag=50)
        assert isinstance(violation, float)
        assert violation >= 0.0

    def test_fdt_monitor_stateful(self):
        """FDTMonitor should accumulate basins and compute violation."""
        monitor = FDTMonitor(max_history=300, max_lag=20)

        # Not ready yet
        assert not monitor.ready
        assert monitor.last_violation == 0.0

        # Record spontaneous basins
        for _ in range(100):
            monitor.record_basin(_random_basin())

        # Still not ready — no perturbation responses
        assert not monitor.ready

        # Record perturbation responses
        for _ in range(25):
            monitor.record_perturbation_response(_random_basin())

        # Now ready
        assert monitor.ready

        # Compute violation
        v = monitor.compute_violation()
        assert isinstance(v, float)
        assert v >= 0.0
        assert monitor.last_violation == v

    def test_fdt_monitor_clear_perturbation(self):
        """clear_perturbation() should reset responses."""
        monitor = FDTMonitor(max_lag=10)
        for _ in range(50):
            monitor.record_basin(_random_basin())
        for _ in range(15):
            monitor.record_perturbation_response(_random_basin())
        assert monitor.ready

        monitor.clear_perturbation()
        assert not monitor.ready


# ─── Integration: Branching ratio + PCI together ──────────────────


class TestIntegration:
    """Sanity-check that all new diagnostics work together."""

    def test_full_diagnostic_suite(self):
        """Run all three diagnostics on a single bank/trajectory."""
        bank = _make_bank(300)

        # Generate a trajectory of activations
        history = []
        for _ in range(20):
            basin = _random_basin()
            bank.activate(basin, top_k=32)
            history.append(basin)

        # 1. Branching ratio
        sigma = bank.branching_ratio()
        assert 0.0 <= sigma <= 1.0

        # 2. PCI
        pci = compute_basin_pci(
            bank.activate,
            history[-1],
            perturbation_scale=0.15,
            response_window=8,
            top_k=16,
        )
        assert pci >= 0.0

        # 3. FDT violation (needs perturbation data)
        perturbation_responses = []
        perturbed = to_simplex(0.8 * history[-1] + 0.2 * _random_basin())
        for _ in range(60):
            bank.activate(perturbed, top_k=16)
            perturbation_responses.append(perturbed.copy())
            perturbed = to_simplex(0.95 * perturbed + 0.05 * _random_basin())

        violation = compute_fdt_violation(
            history * 10,  # repeat to get enough samples
            perturbation_responses,
            max_lag=50,
        )
        assert violation >= 0.0
