"""
Tests for optimized kernel and seam accounting implementations.

Validates:
- OPT-1: Homogeneity detection correctness and performance
- OPT-4: Log-space κ computation correctness
- OPT-10: Incremental seam ledger correctness
- OPT-11: Residual accumulation detection

Interconnections:
- Tests: kernel_optimized.py, seam_optimized.py
- Validates: COMPUTATIONAL_OPTIMIZATIONS.md strategies
- Ensures: Numerical equivalence to formal specification
"""

import time

import numpy as np
import pytest
from umcp.kernel_optimized import (
    CoherenceAnalyzer,
    OptimizedKernelComputer,
    ThresholdCalibrator,
)
from umcp.seam_optimized import (
    SeamChainAccumulator,
    SeamCompositionAnalyzer,
    validate_seam_residuals,
)


class TestOptimizedKernelComputer:
    """Test suite for optimized kernel computation."""

    def test_homogeneous_detection(self):
        """OPT-1: Verify homogeneity detection works correctly."""
        computer = OptimizedKernelComputer(epsilon=1e-6)

        # Homogeneous case
        c_homo = np.array([0.5, 0.5, 0.5, 0.5])
        w = np.array([0.25, 0.25, 0.25, 0.25])

        outputs = computer.compute(c_homo, w)

        assert outputs.is_homogeneous
        assert outputs.computation_mode == "fast_homogeneous"
        assert outputs.C == 0.0  # Lemma 10: C = 0 iff homogeneous
        assert abs(outputs.F - outputs.IC) < 1e-10  # Lemma 4: F = IC iff homogeneous
        assert outputs.amgm_gap < 1e-10

    def test_heterogeneous_computation(self):
        """Verify heterogeneous computation is correct."""
        computer = OptimizedKernelComputer(epsilon=1e-6)

        # Heterogeneous case
        c_hetero = np.array([0.3, 0.5, 0.7, 0.9])
        w = np.array([0.25, 0.25, 0.25, 0.25])

        outputs = computer.compute(c_hetero, w)

        assert not outputs.is_homogeneous
        assert outputs.computation_mode == "full_heterogeneous"
        assert outputs.C > 0.0  # Dispersion present
        assert outputs.F > outputs.IC  # Lemma 4: AM-GM inequality

    def test_homogeneous_vs_heterogeneous_equivalence(self):
        """Verify fast path produces same results as full path for homogeneous data."""
        computer = OptimizedKernelComputer(epsilon=1e-6)

        # Create nearly homogeneous data (within tolerance)
        c = np.array([0.5, 0.5, 0.5, 0.5])
        w = np.array([0.25, 0.25, 0.25, 0.25])

        outputs = computer.compute(c, w)

        # Manually compute with heterogeneous formulas
        F_expected = np.sum(w * c)
        kappa_expected = np.sum(w * np.log(c))
        IC_expected = np.exp(kappa_expected)

        assert abs(outputs.F - F_expected) < 1e-10
        assert abs(outputs.kappa - kappa_expected) < 1e-10
        assert abs(outputs.IC - IC_expected) < 1e-10

    def test_range_validation_lemma1(self):
        """OPT-2: Verify Lemma 1 range bounds are enforced."""
        computer = OptimizedKernelComputer(epsilon=1e-6)

        # Valid case
        c = np.array([0.3, 0.5, 0.7])
        w = np.array([1 / 3, 1 / 3, 1 / 3])

        outputs = computer.compute(c, w, validate=True)  # Should not raise
        assert outputs.F >= 0  # Verify outputs are valid

        # Out of range case: F > 1 (impossible with valid inputs, test validation)
        # We can't create invalid inputs naturally, so test the validation directly
        from umcp.kernel_optimized import validate_kernel_bounds

        assert validate_kernel_bounds(0.5, 0.5, 0.2, 0.5, 0.1, epsilon=1e-6)
        assert not validate_kernel_bounds(1.5, 0.5, 0.2, 0.5, 0.1, epsilon=1e-6)  # F > 1
        assert not validate_kernel_bounds(0.5, 0.5, 0.2, 2.0, 0.1, epsilon=1e-6)  # IC > 1

    def test_log_space_computation_lemma2(self):
        """OPT-4: Verify log-space κ computation is more stable."""
        computer = OptimizedKernelComputer(epsilon=1e-6)

        # Small coordinates where exp/log round-trip loses precision
        c = np.array([1e-5, 1e-5, 1e-5])
        w = np.array([1 / 3, 1 / 3, 1 / 3])

        outputs = computer.compute(c, w)

        # Verify IC = exp(κ) relationship (Lemma 2)
        IC_from_kappa = np.exp(outputs.kappa)
        assert abs(outputs.IC - IC_from_kappa) < 1e-10

        # Verify κ was computed correctly
        kappa_expected = np.sum(w * np.log(c))
        assert abs(outputs.kappa - kappa_expected) < 1e-10

    def test_amgm_gap_lemma4_lemma34(self):
        """OPT-3: Verify AM-GM gap computation (Lemmas 4, 34)."""
        computer = OptimizedKernelComputer(epsilon=1e-6)

        # Heterogeneous case: F > IC
        c = np.array([0.2, 0.8])
        w = np.array([0.5, 0.5])

        outputs = computer.compute(c, w)

        # Lemma 4: F >= IC, equality iff homogeneous
        assert outputs.F >= outputs.IC
        assert outputs.amgm_gap >= 0

        # Manual verification
        F_manual = 0.5 * 0.2 + 0.5 * 0.8  # = 0.5
        IC_manual = (0.2**0.5) * (0.8**0.5)  # = 0.4
        gap_manual = F_manual - IC_manual

        assert abs(outputs.amgm_gap - gap_manual) < 1e-10

    def test_lipschitz_error_propagation_lemma23(self):
        """OPT-12: Verify Lipschitz error bounds (Lemma 23)."""
        computer = OptimizedKernelComputer(epsilon=1e-6)

        # Perturb coordinates by δ
        c1 = np.array([0.4, 0.5, 0.6])
        c2 = np.array([0.41, 0.51, 0.61])  # δ = 0.01
        w = np.array([1 / 3, 1 / 3, 1 / 3])

        delta_c = 0.01

        outputs1 = computer.compute(c1, w)
        outputs2 = computer.compute(c2, w)

        # Compute actual errors
        error_F = abs(outputs1.F - outputs2.F)
        error_omega = abs(outputs1.omega - outputs2.omega)
        error_kappa = abs(outputs1.kappa - outputs2.kappa)
        error_S = abs(outputs1.S - outputs2.S)

        # Verify Lemma 23 bounds (with small tolerance for floating point)
        error_bounds = computer.propagate_coordinate_error(delta_c)

        assert error_F <= error_bounds.F * 1.0001  # Allow tiny floating point slack
        assert error_omega <= error_bounds.omega * 1.0001
        assert error_kappa <= error_bounds.kappa * 1.0001
        assert error_S <= error_bounds.S * 1.0001

    def test_performance_homogeneous_detection(self):
        """OPT-1: Demonstrate performance benefit of homogeneity detection."""
        computer = OptimizedKernelComputer(epsilon=1e-6)

        # Homogeneous data
        c_homo = np.ones(100) * 0.5
        w = np.ones(100) / 100

        # Warm-up
        for _ in range(10):
            computer.compute(c_homo, w)

        # Time homogeneous path
        start = time.time()
        for _ in range(1000):
            computer.compute(c_homo, w)
        time_homo = time.time() - start

        # Heterogeneous data (force full computation)
        c_hetero = np.linspace(0.1, 0.9, 100)

        # Time heterogeneous path
        start = time.time()
        for _ in range(1000):
            computer.compute(c_hetero, w)
        time_hetero = time.time() - start

        # Homogeneous should be faster (typically 20-40% improvement)
        speedup = time_hetero / time_homo
        print(f"\nHomogeneous speedup: {speedup:.2f}x")

        # Not asserting specific speedup as it varies by system
        # But we verify it doesn't slow down
        assert speedup >= 1.0


class TestCoherenceAnalyzer:
    """Test suite for coherence proxy (OPT-14, Lemma 26)."""

    def test_coherence_proxy_bounds(self):
        """Verify Θ ∈ [0, 2] (Lemma 26)."""
        # Test extremes
        theta_min = CoherenceAnalyzer.compute_coherence_proxy(omega=1.0, S=0.0)
        theta_max = CoherenceAnalyzer.compute_coherence_proxy(omega=0.0, S=np.log(2))

        assert 0 <= theta_min <= 2
        assert 0 <= theta_max <= 2

        # Lemma 26: Θ(t) = 1 - ω(t) + S(t)/ln(2)
        assert abs(theta_min - 0.0) < 1e-10
        assert abs(theta_max - 2.0) < 1e-10

    def test_coherence_classification(self):
        """Verify coherence regime classification."""
        # Collapse regime
        theta_collapse = 0.3
        assert CoherenceAnalyzer.classify_coherence(theta_collapse) == "COLLAPSE"

        # Marginal regime
        theta_marginal = 0.7
        assert CoherenceAnalyzer.classify_coherence(theta_marginal) == "MARGINAL"

        # Coherent regime
        theta_coherent = 1.5
        assert CoherenceAnalyzer.classify_coherence(theta_coherent) == "COHERENT"


class TestThresholdCalibrator:
    """Test suite for adaptive threshold calibration (OPT-15, Lemma 34)."""

    def test_threshold_calibration(self):
        """Verify drift threshold calibration via AM-GM gap."""
        # Homogeneous case: gap = 0 → no adjustment
        threshold_homo = ThresholdCalibrator.calibrate_omega_threshold(F=0.5, IC=0.5, base_threshold=0.3)
        assert abs(threshold_homo - 0.3) < 1e-10

        # Heterogeneous case: gap > 0 → tighten threshold
        threshold_hetero = ThresholdCalibrator.calibrate_omega_threshold(F=0.6, IC=0.4, base_threshold=0.3)
        # gap = 0.2, adjustment = base * (1 - 2*0.2) = 0.3 * 0.6 = 0.18
        assert threshold_hetero < 0.3
        assert threshold_hetero >= 0.1  # Clipped lower bound


class TestSeamChainAccumulator:
    """Test suite for optimized seam accounting (OPT-10, OPT-11)."""

    def test_incremental_ledger_update_lemma20(self):
        """OPT-10: Verify incremental ledger composition (Lemma 20)."""
        chain = SeamChainAccumulator()

        # Add three seams
        chain.add_seam(t0=0, t1=10, kappa_t0=0.0, kappa_t1=0.1, tau_R=5.0)
        chain.add_seam(t0=10, t1=20, kappa_t0=0.1, kappa_t1=0.25, tau_R=6.0)
        chain.add_seam(t0=20, t1=30, kappa_t0=0.25, kappa_t1=0.4, tau_R=7.0)

        # Lemma 20: Total = sum of individual changes
        expected_total = (0.1 - 0.0) + (0.25 - 0.1) + (0.4 - 0.25)
        assert abs(chain.get_total_change() - expected_total) < 1e-10
        assert abs(chain.get_total_change() - 0.4) < 1e-10

        # Verify O(1) query is correct
        manual_sum = sum(s.delta_kappa_ledger for s in chain.seam_history)
        assert abs(chain.get_total_change() - manual_sum) < 1e-10

    def test_residual_accumulation_detection_lemma27(self):
        """OPT-11: Verify residual growth detection (Lemma 27)."""
        # Create chain without auto-failure for testing
        chain = SeamChainAccumulator(alpha=0.05, K_max=100)

        # Add seams with controlled residuals (sublinear growth)
        np.random.seed(42)
        for k in range(50):
            # Create budget-balanced seams with small random residuals
            kappa_t0 = k * 0.01
            kappa_t1 = (k + 1) * 0.01  # Clean growth
            tau_R = 10.0

            # Budget model: make R match the ledger change + small noise
            # This ensures residuals are small and zero-mean
            ledger_change = kappa_t1 - kappa_t0  # = 0.01
            R = ledger_change / tau_R  # = 0.001 exactly
            noise = np.random.normal(0, 0.0001)  # Very small noise

            try:
                chain.add_seam(
                    t0=k * 10,
                    t1=(k + 1) * 10,
                    kappa_t0=kappa_t0,
                    kappa_t1=kappa_t1,
                    tau_R=tau_R,
                    R=R + noise,  # Add noise to R, not to ledger
                )
            except ValueError:
                # Growth detection may trigger - check if it's premature
                pass

        # Check final metrics - should show returning dynamics
        metrics = chain.get_metrics()

        # Random walk growth: expect exponent around 0.5-1.2 depending on sample
        # The key is it's not strongly superlinear (> 1.5)
        assert metrics.growth_exponent < 1.5, f"Growth exponent {metrics.growth_exponent:.3f} too high (>1.5)"
        # Mean residual should be small
        assert metrics.mean_residual < 0.01

    def test_residual_failure_detection(self):
        """OPT-11: Verify failure detection for linear growth."""
        chain = SeamChainAccumulator(alpha=0.05)

        # Add seams with linearly growing residuals (model failure)
        try:
            for k in range(30):
                kappa_t0 = k * 0.01
                kappa_t1 = (k + 1) * 0.01

                # Residual grows linearly: s_k = 0.01 * k
                # This simulates a bad budget model
                chain.add_seam(
                    t0=k * 10,
                    t1=(k + 1) * 10,
                    kappa_t0=kappa_t0,
                    kappa_t1=kappa_t1,
                    tau_R=0.0,  # No return credit
                    R=0.0,  # No budget
                    D_omega=0.01 * k,  # Growing penalty
                )

            # Should detect failure before reaching 30 seams
            raise AssertionError("Should have raised ValueError")

        except ValueError as e:
            assert "non-returning dynamics" in str(e)

    def test_seam_composition_validation(self):
        """Verify Lemma 20 composition law validation."""
        chain = SeamChainAccumulator()

        chain.add_seam(t0=0, t1=10, kappa_t0=0.0, kappa_t1=0.1, tau_R=5.0)
        chain.add_seam(t0=10, t1=20, kappa_t0=0.1, kappa_t1=0.25, tau_R=6.0)
        chain.add_seam(t0=20, t1=30, kappa_t0=0.25, kappa_t1=0.4, tau_R=7.0)

        # Validate composition
        result = SeamCompositionAnalyzer.validate_composition_law(chain, t0_chain=0, t2_chain=30)

        assert result["valid"]
        assert abs(result["composed_total"] - result["direct_total"]) < 1e-10
        assert result["num_seams"] == 3

    def test_residual_pattern_analysis(self):
        """Test statistical analysis of residual patterns."""
        # Well-behaved residuals (zero-mean, bounded)
        # Need larger sample (200) for stable growth exponent estimation
        np.random.seed(123)
        residuals_good = list(np.random.normal(0, 0.01, 200))

        analysis = SeamCompositionAnalyzer.analyze_residual_pattern(residuals_good)

        assert analysis["valid"]
        # With 200 samples, random walk should show growth < 1.2
        assert analysis["growth_exponent"] < 1.2
        assert analysis["is_centered"]  # Should be centered around zero

    def test_performance_incremental_vs_recomputation(self):
        """OPT-10: Demonstrate performance of incremental updates."""
        # Incremental approach
        chain_incremental = SeamChainAccumulator()

        start = time.time()
        for k in range(1000):
            chain_incremental.add_seam(
                t0=k * 10,
                t1=(k + 1) * 10,
                kappa_t0=k * 0.01,
                kappa_t1=(k + 1) * 0.01,
                tau_R=10.0,
            )
        time_incremental = time.time() - start
        assert time_incremental > 0  # Sanity check

        # Verify O(1) query
        start = time.time()
        totals = [chain_incremental.get_total_change() for _ in range(1000)]
        time_query = time.time() - start
        assert len(totals) == 1000  # Verify execution

        # Recomputation approach (O(K) every time)
        start = time.time()
        totals_recomp = [sum(s.delta_kappa_ledger for s in chain_incremental.seam_history) for _ in range(1000)]
        time_recompute = time.time() - start
        assert len(totals_recomp) == 1000  # Verify execution

        print(f"\nIncremental query: {time_query:.6f}s")
        print(f"Recomputation: {time_recompute:.6f}s")
        print(f"Speedup: {time_recompute / time_query:.2f}x")

        # Incremental should be faster
        assert time_query < time_recompute


class TestValidationFunctions:
    """Test convenience validation functions."""

    def test_seam_residual_validation(self):
        """Test residual sequence validation function."""
        # Good residuals
        residuals_good = list(np.random.normal(0, 0.01, 50))
        assert validate_seam_residuals(residuals_good)

        # Bad residuals (linear growth)
        residuals_bad = list(np.arange(50) * 0.1)
        assert not validate_seam_residuals(residuals_bad)


class TestComputeUtils:
    """Test suite for compute_utils.py utilities."""

    def test_prune_zero_weights(self):
        """OPT-17: Verify zero-weight dimension pruning."""
        from umcp.compute_utils import prune_zero_weights

        c = np.array([0.3, 0.5, 0.7, 0.9])
        w = np.array([0.5, 0.0, 0.5, 0.0])

        result = prune_zero_weights(c, w)

        assert result.n_original == 4
        assert result.n_active == 2
        assert len(result.c_active) == 2
        assert result.pruned_indices == [1, 3]
        assert np.allclose(result.w_active, [0.5, 0.5])  # Renormalized

    def test_prune_all_active(self):
        """Test pruning with no zeros."""
        from umcp.compute_utils import prune_zero_weights

        c = np.array([0.3, 0.5, 0.7])
        w = np.array([0.4, 0.3, 0.3])

        result = prune_zero_weights(c, w)

        assert result.n_active == result.n_original
        assert len(result.pruned_indices) == 0
        assert np.array_equal(result.c_active, c)

    def test_clip_coordinates(self):
        """Test coordinate clipping with diagnostics."""
        from umcp.compute_utils import clip_coordinates

        c = np.array([0.0, 0.5, 1.0, 0.3])
        result = clip_coordinates(c, epsilon=0.01)

        assert result.clip_count == 2  # 0.0 and 1.0
        assert 0 in result.oor_indices
        assert 2 in result.oor_indices
        assert np.all(result.c_clipped >= 0.01)
        assert np.all(result.c_clipped <= 0.99)

    def test_validate_inputs_valid(self):
        """Test input validation with valid data."""
        from umcp.compute_utils import validate_inputs

        c = np.array([0.3, 0.5, 0.7])
        w = np.array([1 / 3, 1 / 3, 1 / 3])

        result = validate_inputs(c, w)
        assert result["valid"]
        assert result["errors"] == ""

    def test_validate_inputs_invalid(self):
        """Test input validation catches errors."""
        from umcp.compute_utils import validate_inputs

        c = np.array([0.3, 0.5, 0.7])
        w = np.array([0.5, 0.5, 0.5])  # Sums to 1.5

        result = validate_inputs(c, w)
        assert not result["valid"]
        assert "sum to 1.0" in str(result["errors"])

    def test_preprocess_trace_row(self):
        """Test complete row preprocessing."""
        from umcp.compute_utils import preprocess_trace_row

        c = np.array([0.0, 0.5, 1.0, 0.3])
        w = np.array([0.25, 0.0, 0.25, 0.5])

        c_proc, w_proc, diagnostics = preprocess_trace_row(c, w, epsilon=0.01)

        # Should have pruned 1 dimension and clipped coordinates
        assert len(c_proc) == 3  # One zero-weight pruned
        assert diagnostics["pruning"] is not None  # Pruning was applied
        assert np.all(c_proc >= 0.01)
        assert np.all(c_proc <= 0.99)
        assert np.allclose(w_proc.sum(), 1.0)


class TestReturnTimeOptimizations:
    """Test suite for tau_R_optimized.py return time optimizations."""

    @pytest.fixture
    def return_computer(self):
        """Load and instantiate the OptimizedReturnComputer."""
        import importlib.util
        import sys
        from pathlib import Path

        # Use relative path from test file location
        repo_root = Path(__file__).parent.parent
        tau_r_path = repo_root / "closures" / "tau_R_optimized.py"

        spec = importlib.util.spec_from_file_location("tau_R_optimized", str(tau_r_path))
        assert spec is not None and spec.loader is not None
        module = importlib.util.module_from_spec(spec)
        # Register module in sys.modules before exec to allow dataclass decorator
        sys.modules["tau_R_optimized"] = module
        spec.loader.exec_module(module)
        return module.OptimizedReturnComputer(H_rec=100, eta=0.1)

    def test_early_exit_with_margin(self, return_computer):  # type: ignore[no-untyped-def]
        """OPT-7: Verify margin-based early exit."""
        # Create a trace where current state matches a past state exactly
        trace = np.array(
            [
                [0.5, 0.5, 0.5],  # t=0
                [0.6, 0.5, 0.5],  # t=1 - moved
                [0.5, 0.5, 0.5],  # t=2 - back to origin
            ]
        )
        psi_t = np.array([0.5, 0.5, 0.5])

        result = return_computer.compute_tau_R(psi_t, trace, t=2)  # type: ignore[misc]

        assert result.tau_R == 2.0  # Returns to t=0  # type: ignore[misc]
        assert result.computation_mode in ["early_exit", "full_search"]  # type: ignore[misc]

    def test_coverage_caching(self, return_computer):  # type: ignore[no-untyped-def]
        """OPT-8: Verify coverage set caching."""
        # Create a simple trace
        trace = np.array(
            [
                [0.5, 0.5, 0.5],
                [0.5, 0.5, 0.5],
                [0.5, 0.5, 0.5],
            ]
        )

        # Get coverage - should cache result
        coverage = return_computer.get_coverage_set(trace, t=2)  # type: ignore[misc]
        coverage2 = return_computer.get_coverage_set(trace, t=2)  # type: ignore[misc]

        assert coverage == coverage2
        # Cache should have entry - hasattr is enough to verify caching works
        assert hasattr(return_computer, "_coverage_cache")  # type: ignore[misc]

    def test_binary_search_eta(self, return_computer):  # type: ignore[no-untyped-def]
        """OPT-9: Verify binary search for minimal η via find_minimal_eta."""
        # Create a trajectory that returns
        trajectory = np.array(
            [
                [0.5, 0.5, 0.5],  # Start at center
                [0.6, 0.5, 0.5],  # Move away
                [0.7, 0.5, 0.5],  # Further away
                [0.55, 0.5, 0.5],  # Start returning
                [0.5, 0.5, 0.5],  # Back at center
            ]
        )
        psi_t = trajectory[4]  # Current state

        # Use the find_minimal_eta method
        eta_min = return_computer.find_minimal_eta(  # type: ignore[misc]
            psi_t, trajectory, t=4
        )

        # Should find an eta where trajectory returns (exact match to t=0)
        assert eta_min is not None
        assert eta_min >= 0.0
        assert eta_min <= 1.0


if __name__ == "__main__":
    pytest.main([__file__, "-v", "-s"])
