"""Tests for the Watkins Kaleidoscope Tensor (3x3x2 transfer operator)."""

import numpy as np
from math import factorial

from watkins_nn.tensor import KaleidoscopeTensor
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():
    for n in [0, 1, 5, 10]:
        assert T.tensor_at(n).shape == (3, 3, 2)

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))


# -- 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] == 3.0

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
