"""Tests for Theorem R: Coupling Energy as Off-Diagonal Interaction,
and the Coupling Angle invariant.

Key identity: E(n) = [tr(M²) - Σᵢ M[i,i]²] / 2 = Σ_{i<j} M[i,j]·M[j,i]

The coupling angle θ(n) measures alignment between upper and lower
triangular halves of M(n).
"""

import math
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.tensor import KaleidoscopeTensor
from watkins_nn.coupling import (
    coupling_energy,
    upper_triangle,
    lower_triangle,
    coupling_angle,
    coupling_torque,
    gram_matrix,
    trace_squared_decomposition,
)

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


# ── Theorem R: E(n) = off-diagonal Frobenius / 2 ────────────────────────────

def test_coupling_equals_off_diagonal():
    """E(n) = Σ_{i<j} M[i,j]·M[j,i] for n=0..50."""
    for n in range(51):
        a = _mat(n)
        off_diag = (a[0][1]*a[1][0] + a[0][2]*a[2][0] + a[1][2]*a[2][1])
        assert coupling_energy(n) == off_diag, f"E({n}) ≠ off-diagonal sum"


def test_coupling_equals_frobenius_formula():
    """E(n) = [tr(M²) - Σᵢ d²ᵢ] / 2 for n=0..50 (exact integer)."""
    for n in range(51):
        a = _mat(n)
        tr_sq = sum(a[i][j] * a[j][i] for i in range(3) for j in range(3))
        diag_sq = sum(a[i][i]**2 for i in range(3))
        assert coupling_energy(n) == (tr_sq - diag_sq) // 2, \
            f"Frobenius formula failed at n={n}"


def test_coupling_positive():
    """E(n) ≥ 1 for all n ≥ 0."""
    for n in range(100):
        assert coupling_energy(n) >= 1, f"E({n}) < 1"


def test_lucas_identity_from_theorem_P():
    """Theorem P implies L(2n)² = L(4n) + 2 (Lucas duplication)."""
    for n in range(1, 40):
        assert lucas(2*n)**2 == lucas(4*n) + 2, f"Lucas identity failed at n={n}"


# ── Triangle vectors ─────────────────────────────────────────────────────────

def test_upper_lower_match_matrix():
    """Upper/lower triangle vectors match M(n) entries for n=0..30."""
    for n in range(31):
        a = _mat(n)
        u = upper_triangle(n)
        v = lower_triangle(n)
        # Upper: M[0,1]=seq3, M[0,2]=seq1, M[1,2]=seq7
        assert u == (a[0][1], a[0][2], a[1][2]), f"upper mismatch at n={n}"
        # Lower: M[1,0]=seq6, M[2,0]=seq4, M[2,1]=seq5
        assert v == (a[1][0], a[2][0], a[2][1]), f"lower mismatch at n={n}"


def test_dot_product_is_coupling():
    """⟨u, v⟩ = E(n) for n=0..50."""
    for n in range(51):
        u = upper_triangle(n)
        v = lower_triangle(n)
        dot = u[0]*v[0] + u[1]*v[1] + u[2]*v[2]
        assert dot == coupling_energy(n), f"dot product mismatch at n={n}"


# ── Coupling angle ───────────────────────────────────────────────────────────

def test_angle_nonnegative():
    """θ(n) ≥ 0 for n=0..100."""
    for n in range(101):
        assert coupling_angle(n) >= 0, f"θ({n}) < 0"


def test_angle_bounded_by_pi():
    """θ(n) ≤ π for n=0..100."""
    for n in range(101):
        assert coupling_angle(n) <= math.pi + 1e-10, f"θ({n}) > π"


def test_angle_decreasing_trend():
    """θ(n) generally decreases for large n (u and v become parallel)."""
    angles = [coupling_angle(n) for n in range(20, 60)]
    # At least the last angle should be smaller than the first
    assert angles[-1] < angles[0], \
        f"θ not decreasing: θ(20)={angles[0]:.6f}, θ(59)={angles[-1]:.6f}"


# ── Coupling torque ──────────────────────────────────────────────────────────

def test_torque_equals_cross_norm():
    """Torque = |u × v| matches sin(θ)·|u|·|v| for n=1..20 (float-safe)."""
    for n in range(1, 21):
        tau = coupling_torque(n)
        theta = coupling_angle(n)
        u = upper_triangle(n)
        v = lower_triangle(n)
        nu = math.sqrt(sum(x**2 for x in u))
        nv = math.sqrt(sum(x**2 for x in v))
        expected = nu * nv * math.sin(theta)
        assert abs(tau - expected) < max(1e-4, tau * 1e-6), \
            f"torque mismatch at n={n}: {tau} vs {expected}"


# ── Gram matrix ──────────────────────────────────────────────────────────────

def test_gram_determinant_nonneg():
    """det(G) = |u|²|v|² - E² ≥ 0 (Cauchy-Schwarz) for n=0..100."""
    for n in range(101):
        _, _, _, det_G = gram_matrix(n)
        assert det_G >= -1, f"Cauchy-Schwarz violated at n={n}: det(G) = {det_G}"


def test_gram_determinant_is_torque_squared():
    """det(G) = (coupling torque)² for n=0..50."""
    for n in range(51):
        _, _, _, det_G = gram_matrix(n)
        tau = coupling_torque(n)
        assert abs(det_G - tau**2) < max(1e-6, abs(det_G) * 1e-8), \
            f"det(G) ≠ τ² at n={n}"


# ── Frobenius decomposition ──────────────────────────────────────────────────

def test_trace_squared_decomposition():
    """tr(M²) = diag_sq + 2E for n=0..50 (exact integer)."""
    for n in range(51):
        a = _mat(n)
        tr_sq = sum(a[i][j] * a[j][i] for i in range(3) for j in range(3))
        result = trace_squared_decomposition(n)
        assert result['trace_sq'] == tr_sq, \
            f"tr(M²) mismatch at n={n}: {result['trace_sq']} vs {tr_sq}"
        assert result['trace_sq'] == result['diag_sq'] + result['coupling']


# ── Theorem N in angle form ──────────────────────────────────────────────────

def test_theorem_N_angle_form():
    """p₁² - 3p₂ = ½Σ(dᵢ-dⱼ)² + 3E (integer) for n=0..50."""
    for n in range(51):
        a = _mat(n)
        d0, d1, d2 = a[0][0], a[1][1], a[2][2]
        p1 = d0 + d1 + d2

        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]
        p2 = m00 + m11 + m22

        lhs = p1**2 - 3*p2
        diag_sep = ((d0-d1)**2 + (d0-d2)**2 + (d1-d2)**2) // 2
        E = coupling_energy(n)

        assert lhs == diag_sep + 3*E, f"Thm N failed at n={n}: {lhs} vs {diag_sep + 3*E}"
