"""Tests for Theorem Q: Characteristic Polynomial Closed Form.

Verifies that the char poly coefficients (p₁, p₂, p₃) computed from the
9 sequences match the standard linear algebra definitions, and that
Cayley-Hamilton holds: M³ - p₁M² + p₂M - p₃I = 0.
"""

import numpy as np

from watkins_nn.kaleidoscope import (
    lucas,
    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,
)
from watkins_nn.charpoly import (
    charpoly_coefficients,
    charpoly_eval,
    trace_M,
    trace_adj,
    det_M,
    cofactor_diagonal,
)
from watkins_nn.tensor import KaleidoscopeTensor

KT = KaleidoscopeTensor()


def _mat(n):
    """Build M(n) as exact Python ints (no numpy overflow)."""
    s2 = seq2_mersenne_numerator(n)
    s3 = seq3_mersenne_gcd(n)
    s1 = seq1_cycle_lb_numerator(n)
    s6 = seq6_mersenne_lucas_gcd(n)
    s8 = seq8_cross_gcd(n)
    s7 = seq7_binary_preservation(n)
    s4 = seq4_lucas_numerator(n)
    s5 = seq5_lucas_gcd(n)
    l2n = seq9_weight_moments(n)
    return [[s2, s3, s1],
            [s6, s8, s7],
            [s4, s5, l2n]]


def _det3(a):
    return (a[0][0] * (a[1][1] * a[2][2] - a[1][2] * a[2][1])
          - a[0][1] * (a[1][0] * a[2][2] - a[1][2] * a[2][0])
          + a[0][2] * (a[1][0] * a[2][1] - a[1][1] * a[2][0]))


# ── p₁ = trace ────────────────────────────────────────────────────────────────

def test_p1_is_trace():
    """p₁ matches integer trace for n=0..50."""
    for n in range(51):
        a = _mat(n)
        assert trace_M(n) == a[0][0] + a[1][1] + a[2][2], f"p₁ mismatch at n={n}"


# ── p₂ = trace of adjugate ───────────────────────────────────────────────────

def test_p2_is_sum_of_2x2_minors():
    """p₂ equals the sum of principal 2x2 minors for n=0..50."""
    for n in range(51):
        a = _mat(n)
        m00 = a[1][1]*a[2][2] - a[1][2]*a[2][1]
        m11 = a[0][0]*a[2][2] - a[0][2]*a[2][0]
        m22 = a[0][0]*a[1][1] - a[0][1]*a[1][0]
        assert trace_adj(n) == m00 + m11 + m22, f"p₂ mismatch at n={n}"


def test_cofactor_diagonal_entries():
    """Cofactor diagonal entries match integer 2x2 minors for n=0..50."""
    for n in range(51):
        c00, c11, c22 = cofactor_diagonal(n)
        a = _mat(n)
        # C₀₀ = det of minor (rows 1,2, cols 1,2)
        m00 = a[1][1]*a[2][2] - a[1][2]*a[2][1]
        # C₁₁ = det of minor (rows 0,2, cols 0,2)
        m11 = a[0][0]*a[2][2] - a[0][2]*a[2][0]
        # C₂₂ = det of minor (rows 0,1, cols 0,1)
        m22 = a[0][0]*a[1][1] - a[0][1]*a[1][0]
        assert c00 == m00, f"C₀₀ mismatch at n={n}"
        assert c11 == m11, f"C₁₁ mismatch at n={n}"
        assert c22 == m22, f"C₂₂ mismatch at n={n}"


# ── p₃ = determinant ─────────────────────────────────────────────────────────

def test_p3_is_determinant():
    """p₃ matches integer Sarrus determinant for n=0..50."""
    for n in range(51):
        a = _mat(n)
        assert det_M(n) == _det3(a), f"det mismatch at n={n}"


# ── Cayley-Hamilton ──────────────────────────────────────────────────────────

def test_cayley_hamilton():
    """M³ - p₁M² + p₂M - p₃I = 0 (exact integer) for n=0..30."""
    for n in range(31):
        a = _mat(n)
        p1, p2, p3 = charpoly_coefficients(n)
        # Integer matrix multiply
        def mmul(A, B):
            return [[sum(A[i][k]*B[k][j] for k in range(3)) for j in range(3)] for i in range(3)]
        M2 = mmul(a, a)
        M3 = mmul(M2, a)
        for i in range(3):
            for j in range(3):
                eye_ij = 1 if i == j else 0
                val = M3[i][j] - p1*M2[i][j] + p2*a[i][j] - p3*eye_ij
                assert val == 0, f"Cayley-Hamilton failed at n={n}, entry ({i},{j})={val}"


# ── Eigenvalue product/sum identities ────────────────────────────────────────

def test_eigenvalue_sum_is_p1():
    """Sum of eigenvalues ≈ p₁ for n=1..20 (float-safe range)."""
    for n in range(1, 21):
        eigs = KT.eigenvalues(n)
        p1, _, _ = charpoly_coefficients(n)
        assert abs(sum(eigs) - p1) < max(1e-4, abs(p1) * 1e-10), \
            f"eigenvalue sum mismatch at n={n}"


def test_eigenvalue_product_is_p3():
    """Product of eigenvalues = p₃ = det for n=1..30."""
    for n in range(1, 31):
        eigs = KT.eigenvalues(n)
        _, _, p3 = charpoly_coefficients(n)
        prod = eigs[0] * eigs[1] * eigs[2]
        assert abs(prod - p3) < max(1e-4, abs(p3) * 1e-8), \
            f"eigenvalue product mismatch at n={n}: {prod} vs {p3}"


def test_charpoly_eval_at_eigenvalues():
    """p_n(λᵢ) ≈ 0 for each eigenvalue, n=1..15 (float-safe range)."""
    for n in range(1, 16):
        eigs = KT.eigenvalues(n)
        for lam in eigs:
            val = charpoly_eval(n, lam)
            assert abs(val) < max(0.1, abs(lam)**3 * 1e-5), \
                f"charpoly({n}, {lam}) = {val} ≠ 0"
