import pytest
"""Tests for the Watkins Kaleidoscope Tensor (3x3x4 transfer operator)."""

import numpy as np
from math import factorial

from watkins_nn.tensor import KaleidoscopeTensor, _v_adic, _big_omega
from watkins_nn.constants import PHI_SQUARED
from watkins_nn.kaleidoscope import (
    seq1_cycle_lb_numerator, seq2_mersenne_numerator, seq3_mersenne_gcd,
    seq4_lucas_numerator, seq5_lucas_gcd, seq6_mersenne_lucas_gcd,
    seq7_binary_preservation, seq8_cross_gcd, seq9_weight_moments,
    lucas,
)

T = KaleidoscopeTensor()


# -- Tensor shape and type ---------------------------------------------------

def test_tensor_at_shape_default():
    for n in [0, 1, 5, 10]:
        assert T.tensor_at(n).shape == (3, 3, 4)

def test_tensor_at_shape_3modes():
    for n in [0, 1, 5, 10]:
        assert T.tensor_at(n, n_modes=3).shape == (3, 3, 3)

def test_tensor_at_dtype():
    assert T.tensor_at(0).dtype == np.float64

def test_integer_matrix_matches_mode0():
    for n in [0, 3, 7]:
        tensor = T.tensor_at(n)
        np.testing.assert_array_equal(tensor[:, :, 0], T.integer_matrix(n))

def test_asymptotic_matrix_matches_mode1():
    for n in [0, 3, 7]:
        tensor = T.tensor_at(n)
        np.testing.assert_array_equal(tensor[:, :, 1], T.asymptotic_matrix(n))

def test_arithmetic_depth_matches_mode2():
    for n in [0, 3, 7]:
        tensor = T.tensor_at(n)
        np.testing.assert_array_equal(tensor[:, :, 2], T.arithmetic_depth_matrix(n))


# -- Mode 0 correctness -----------------------------------------------------

def test_mode0_n0():
    M = T.integer_matrix(0)
    expected = np.array([
        [3, 1, 1],   # seq2(0)=3, seq3(0)=1, seq1(0)=1
        [2, 1, 0],   # seq6(0)=2, seq8(0)=1, seq7(0)=0
        [1, 1, 2],   # seq4(0)=1, seq5(0)=1, seq9(0)=L(0)=2
    ], dtype=float)
    np.testing.assert_array_equal(M, expected)

def test_mode0_n1():
    M = T.integer_matrix(1)
    expected = np.array([
        [7, 1, 3],   # seq2(1)=7, seq3(1)=1, seq1(1)=3
        [1, 3, 0],   # seq6(1)=1, seq8(1)=3, seq7(1)=0
        [3, 1, 3],   # seq4(1)=3, seq5(1)=1, seq9(1)=L(2)=3
    ], dtype=float)
    np.testing.assert_array_equal(M, expected)

def test_mode0_n14_resonance():
    M = T.integer_matrix(14)
    assert int(M[1, 1]) == 31  # Cross-GCD spike
    assert int(M[2, 0]) == 31  # Lucas numerator echo


# -- Mode 1 correctness -----------------------------------------------------

def test_mode1_lagrange_burmann_sign():
    M0 = T.asymptotic_matrix(0)
    M1 = T.asymptotic_matrix(1)
    # n=0: sign = (-1)^1 = -1 -> all col 0 values negative (if raw > 0)
    assert M0[0, 0] < 0   # -seq2(0)/1! = -3
    assert M0[2, 0] < 0   # -seq4(0)/1! = -1
    # n=1: sign = (-1)^2 = +1 -> all col 0 values positive
    assert M1[0, 0] > 0   # +seq2(1)/2! = 7/2

def test_mode1_lagrange_burmann_decay():
    """LB coefficients decay factorially."""
    vals = [abs(T.asymptotic_matrix(n)[0, 0]) for n in range(10)]
    # After initial terms, should generally decrease
    for i in range(3, len(vals)):
        assert vals[i] < vals[1], f"LB decay failed at n={i}"

def test_mode1_resonance_zero_when_trivial():
    """At n=0, all GCD column values equal 1 -> Mode 1 col 1 all zeros."""
    M = T.asymptotic_matrix(0)
    np.testing.assert_array_equal(M[:, 1], [0.0, 0.0, 0.0])

def test_mode1_resonance_nonzero_at_spike():
    """At n=1, seq8(1)=3 -> Mode 1 cell (1,1) = 3."""
    M = T.asymptotic_matrix(1)
    assert M[1, 1] == pytest.approx(3.0, abs=1e-15)

def test_mode1_phi_scaled_seq9_converges():
    """seq9(k)/phi^(2k) = L(2k)/phi^(2k) -> 1.0 as k -> inf."""
    val = seq9_weight_moments(20) / (PHI_SQUARED ** 20)
    assert abs(val - 1.0) < 1e-8


# -- Derived quantities -----------------------------------------------------

def test_curvature_n0():
    # trace = 3+1+2 = 6, det = 3*(1*2-0*1) - 1*(2*2-0*1) + 1*(2*1-1*1)
    # = 3*2 - 1*4 + 1*1 = 6 - 4 + 1 = 3
    assert abs(T.curvature(0) - 2.0) < 1e-10

def test_curvature_handles_singular():
    """If det is 0, curvature returns inf."""
    # We test the logic by checking a known non-singular case
    curv = T.curvature(0)
    assert curv != float('inf')
    assert abs(curv - 2.0) < 1e-10

def test_spectral_gap_positive():
    for n in range(20):
        assert T.spectral_gap(n) >= 0

def test_eigenvalues_sorted_descending():
    for n in [0, 1, 5, 14]:
        eigs = T.eigenvalues(n)
        assert len(eigs) == 3
        for i in range(len(eigs) - 1):
            assert eigs[i] >= eigs[i + 1]


# -- Resonance slice --------------------------------------------------------

def test_resonance_slice_keys():
    res = T.resonance_slice(0)
    expected_keys = {
        'n', 'kappa_channel', 'kappa_spikes', 'resonant_rows',
        'cross_gcd', 'is_bridge_resonance', 'curvature',
        'spectral_gap', 'eigenvalues', 'lagrange_burmann_col',
        'phi_scaled_col',
    }
    assert set(res.keys()) == expected_keys

def test_resonance_slice_bridge_at_14():
    res = T.resonance_slice(14)
    assert res['is_bridge_resonance'] is True
    assert res['cross_gcd'] == 31


# -- Print -------------------------------------------------------------------

def test_print_tensor_returns_string(capsys):
    out = T.print_tensor(0)
    assert isinstance(out, str)
    assert "KALEIDOSCOPE TENSOR" in out
    assert "Mode 0" in out
    assert "Mode 1" in out
    assert "Mode 2" in out


# -- v-adic and Omega helpers ------------------------------------------------

def test_v_adic_basic():
    assert _v_adic(0, 2) == 0
    assert _v_adic(1, 2) == 0
    assert _v_adic(2, 2) == 1
    assert _v_adic(8, 2) == 3
    assert _v_adic(12, 2) == 2
    assert _v_adic(12, 3) == 1

def test_v_adic_5():
    assert _v_adic(25, 5) == 2
    assert _v_adic(125, 5) == 3
    assert _v_adic(7, 5) == 0

def test_big_omega_basic():
    assert _big_omega(0) == 0
    assert _big_omega(1) == 0
    assert _big_omega(2) == 1
    assert _big_omega(12) == 3   # 2^2 * 3
    assert _big_omega(30) == 3   # 2 * 3 * 5


# -- Mode 2 correctness -----------------------------------------------------

def test_mode2_col0_is_v2():
    """Col 0 of Mode 2 = v_2(val) for all entries."""
    for n in [0, 1, 5, 10, 14]:
        M0 = T.integer_matrix(n)
        M2 = T.arithmetic_depth_matrix(n)
        for r in range(3):
            assert M2[r, 0] == _v_adic(int(M0[r, 0]), 2)

def test_mode2_col1_is_v5():
    """Col 1 of Mode 2 = v_5(val) for all entries."""
    for n in [0, 1, 5, 10, 14]:
        M0 = T.integer_matrix(n)
        M2 = T.arithmetic_depth_matrix(n)
        for r in range(3):
            assert M2[r, 1] == _v_adic(int(M0[r, 1]), 5)

def test_mode2_col2_is_omega():
    """Col 2 of Mode 2 = Omega(val) for val > 1, else 0."""
    for n in [1, 5, 10, 14]:
        M0 = T.integer_matrix(n)
        M2 = T.arithmetic_depth_matrix(n)
        for r in range(3):
            val = abs(int(M0[r, 2]))
            expected = _big_omega(val) if val > 1 else 0
            assert M2[r, 2] == expected

def test_mode2_n0_values():
    """Mode 2 at n=0 has known values."""
    M2 = T.arithmetic_depth_matrix(0)
    # seq2(0)=3 -> v_2(3)=0, seq3(0)=1 -> v_5(1)=0, seq1(0)=1 -> Omega=0
    assert M2[0, 0] == 0  # v_2(3)
    assert M2[0, 1] == 0  # v_5(1)
    assert M2[0, 2] == 0  # Omega(1)=0
    # seq6(0)=2 -> v_2(2)=1
    assert M2[1, 0] == 1  # v_2(2)
    # seq9(0)=L(0)=2 -> Omega(2)=1
    assert M2[2, 2] == 1  # Omega(2)

def test_mode2_mersenne_no_2adic():
    """Mersenne numbers 2^n-1 are always odd -> v_2 = 0 for seq2."""
    for n in range(1, 20):
        M2 = T.arithmetic_depth_matrix(n)
        assert M2[0, 0] == 0, f"seq2({n}) has v_2 > 0"

def test_mode2_nonnegative():
    """All Mode 2 entries are non-negative integers."""
    for n in range(20):
        M2 = T.arithmetic_depth_matrix(n)
        for r in range(3):
            for c in range(3):
                assert M2[r, c] >= 0
                assert M2[r, c] == int(M2[r, c])

def test_mode2_deterministic():
    """Mode 2 gives identical results on repeated calls."""
    for n in [0, 5, 14]:
        a = T.arithmetic_depth_matrix(n)
        b = T.arithmetic_depth_matrix(n)
        np.testing.assert_array_equal(a, b)


# -- Mode 3 correctness (cofactor interaction) --------------------------------

def test_cofactor_matrix_matches_mode3():
    for n in [0, 3, 7, 14]:
        tensor = T.tensor_at(n)
        np.testing.assert_array_equal(tensor[:, :, 3], T.cofactor_matrix(n))

def test_cofactor_adjugate_identity():
    """M * adj(M) = det(M) * I, where adj = cofactor^T."""
    for n in [1, 3, 5, 8, 14]:
        M0 = T.integer_matrix(n)
        C = T.cofactor_matrix(n)
        adj = C.T  # adjugate = transpose of cofactor
        product = M0 @ adj
        det_M = np.linalg.det(M0)
        expected = det_M * np.eye(3)
        np.testing.assert_allclose(product, expected, atol=1e-6)

def test_cofactor_c11_is_delta_ml():
    """C[1,1] = seq2*L(2n) - seq1*seq4 (Mersenne-Lucas minor, Theorem I)."""
    for n in [1, 5, 10, 14, 20]:
        C = T.cofactor_matrix(n)
        s2 = seq2_mersenne_numerator(n)
        s1 = seq1_cycle_lb_numerator(n)
        s4 = seq4_lucas_numerator(n)
        L2n = seq9_weight_moments(n)
        delta_ml = s2 * L2n - s1 * s4
        assert abs(C[1, 1] - delta_ml) < 1e-6, f"C[1,1] mismatch at n={n}"

def test_cofactor_c00_formula():
    """C[0,0] = seq8*L(2n) - seq7*seq5."""
    for n in [1, 5, 10, 14]:
        C = T.cofactor_matrix(n)
        s8 = seq8_cross_gcd(n)
        s7 = seq7_binary_preservation(n)
        s5 = seq5_lucas_gcd(n)
        L2n = seq9_weight_moments(n)
        expected = s8 * L2n - s7 * s5
        assert abs(C[0, 0] - expected) < 1e-6

def test_cofactor_c22_formula():
    """C[2,2] = seq2*seq8 - seq3*seq6."""
    for n in [1, 5, 10, 14]:
        C = T.cofactor_matrix(n)
        s2 = seq2_mersenne_numerator(n)
        s3 = seq3_mersenne_gcd(n)
        s6 = seq6_mersenne_lucas_gcd(n)
        s8 = seq8_cross_gcd(n)
        expected = s2 * s8 - s3 * s6
        assert abs(C[2, 2] - expected) < 1e-6

def test_cofactor_trace_is_i2():
    """tr(cofactor) = I₂ (second invariant of characteristic polynomial)."""
    for n in [1, 3, 5, 10]:
        M0 = T.integer_matrix(n)
        C = T.cofactor_matrix(n)
        trace_cof = C[0,0] + C[1,1] + C[2,2]
        # I₂ = sum of 2×2 principal minors
        i2 = (M0[0,0]*M0[1,1] - M0[0,1]*M0[1,0] +
              M0[0,0]*M0[2,2] - M0[0,2]*M0[2,0] +
              M0[1,1]*M0[2,2] - M0[1,2]*M0[2,1])
        assert abs(trace_cof - i2) < 1e-6

def test_cofactor_deterministic():
    """Mode 3 gives identical results on repeated calls."""
    for n in [0, 5, 14]:
        a = T.cofactor_matrix(n)
        b = T.cofactor_matrix(n)
        np.testing.assert_array_equal(a, b)

def test_cofactor_integer_valued():
    """All cofactor entries are integers (since M(n) has integer entries)."""
    for n in range(20):
        C = T.cofactor_matrix(n)
        for r in range(3):
            for c in range(3):
                assert abs(C[r, c] - round(C[r, c])) < 1e-10

def test_print_tensor_includes_mode3(capsys):
    out = T.print_tensor(0)
    assert "Mode 3" in out
    assert "cofactor" in out.lower()
    assert "Delta_ML" in out
