"""Tests for Theorem U: Arithmetic Factorization of Coupling Energy.

Theorem U states:

    E(n) = (2φ)^{n+1} / D(n) + O(φ^{-n})

where D(n) = gcd(M(n+1), (n+1)!) × gcd(L(n+1), (n+1)!) decomposes over
primes p ≤ n+1.  C_odd = 2φ/9 corresponds to D = 9 = 3² at positions
where only p = 3 contributes.
"""

import math
from math import gcd, factorial

import numpy as np

from watkins_nn.constants import PHI
from watkins_nn.kaleidoscope import mersenne, lucas
from watkins_nn.coupling import (
    coupling_energy,
    arithmetic_denominator,
    coupling_energy_predicted,
)

TWO_PHI = 2 * PHI


# ── Theorem U: master formula ─────────────────────────────────────────────────

def test_master_formula_relative_error():
    """E(n) = (2φ)^{n+1}/D(n) with relative error < 1e-6 for n=30..100."""
    for n in range(30, 101):
        E = coupling_energy(n)
        E_pred = coupling_energy_predicted(n)
        if E == 0:
            continue
        rel = abs(E - E_pred) / E
        assert rel < 1e-6, \
            f"Relative error {rel:.2e} > 1e-6 at n={n}"


def test_master_formula_improves_with_n():
    """Relative error decreases as n increases."""
    errors = []
    for n in [30, 50, 70, 90]:
        E = coupling_energy(n)
        E_pred = coupling_energy_predicted(n)
        errors.append(abs(E - E_pred) / E if E > 0 else 0)
    assert errors[-1] < errors[0], \
        f"Error not improving: {errors}"


def test_predicted_matches_energy():
    """coupling_energy_predicted(n) ≈ coupling_energy(n) for n=25..60."""
    for n in range(25, 61):
        E = coupling_energy(n)
        E_pred = coupling_energy_predicted(n)
        if E > 0:
            assert abs(E_pred / E - 1) < 1e-6, \
                f"Predicted/actual = {E_pred/E} at n={n}"


# ── Arithmetic denominator ─────────────────────────────────────────────────────

def test_denominator_factorization():
    """D(n) = gcd(M(n+1),(n+1)!) × gcd(L(n+1),(n+1)!) for n=0..80."""
    for n in range(81):
        D, gM, gL = arithmetic_denominator(n)
        assert D == gM * gL, f"D ≠ gM×gL at n={n}"
        Mn1 = mersenne(n + 1)
        Ln1 = lucas(n + 1)
        fn1 = factorial(n + 1)
        assert gM == gcd(Mn1, fn1), f"gM mismatch at n={n}"
        assert gL == gcd(Ln1, fn1), f"gL mismatch at n={n}"


def test_denominator_positive():
    """D(n) ≥ 1 for all n."""
    for n in range(100):
        D, _, _ = arithmetic_denominator(n)
        assert D >= 1


def test_raw_product_converges_to_one():
    """M(n+1)×L(n+1) / (2φ)^{n+1} → 1 for large n."""
    for n in [40, 60, 80]:
        Mn1 = mersenne(n + 1)
        Ln1 = lucas(n + 1)
        raw = Mn1 * Ln1 / TWO_PHI ** (n + 1)
        assert abs(raw - 1.0) < 1e-8, \
            f"Raw product = {raw} at n={n}, expected ~1"


# ── C_odd = 2φ/9 derivation ──────────────────────────────────────────────────

def test_c_odd_from_D_equals_9():
    """At MEDIUM positions, D = 9 = 3² and E/(2φ)^n = 2φ/9."""
    C_ODD = TWO_PHI / 9
    for n in [33, 37, 45, 57, 61, 73]:
        D, gM, gL = arithmetic_denominator(n)
        assert D == 9, f"D = {D} ≠ 9 at n={n}"
        assert gM == 3, f"gM = {gM} ≠ 3 at n={n}"
        assert gL == 3, f"gL = {gL} ≠ 3 at n={n}"
        E = coupling_energy(n)
        ratio = E / TWO_PHI ** n
        assert abs(ratio - C_ODD) / C_ODD < 1e-10, \
            f"ratio = {ratio}, expected {C_ODD} at n={n}"


def test_high_from_D_equals_1():
    """At HIGH positions, D = 1 and E/(2φ)^n = 2φ."""
    for n in [36, 40, 42, 58, 60, 66, 70, 72]:
        D, gM, gL = arithmetic_denominator(n)
        assert D == 1, f"D = {D} ≠ 1 at n={n}"
        E = coupling_energy(n)
        ratio = E / TWO_PHI ** n
        assert abs(ratio - TWO_PHI) / TWO_PHI < 1e-10, \
            f"ratio = {ratio}, expected {TWO_PHI} at n={n}"


def test_nine_is_three_squared():
    """The 9 in C_odd = 2φ/9 is 3_M × 3_L from ord₂(3)=2 and α_L(3)=2."""
    # ord_2(3) = 2: 2² ≡ 1 mod 3
    assert pow(2, 2, 3) == 1
    # α_L(3) = 2: 3 | L(2)
    assert lucas(2) % 3 == 0
    assert lucas(1) % 3 != 0  # entry point is exactly 2


# ── D values at specific positions ────────────────────────────────────────────

def test_D_28_positions():
    """D = 28 = 7_M × 4_L at known positions."""
    for n in [38, 50, 56]:
        D, gM, gL = arithmetic_denominator(n)
        assert D == 28, f"D = {D} at n={n}"
        assert gM == 7, f"gM = {gM} at n={n}"
        assert gL == 4, f"gL = {gL} at n={n}"


def test_D_29_position():
    """D = 29 at n=48 (prime 29 from Lucas only)."""
    D, gM, gL = arithmetic_denominator(48)
    assert D == 29
    assert gM == 1
    assert gL == 29


def test_D_341_position():
    """D = 341 = 11×31 at n=64."""
    D, gM, gL = arithmetic_denominator(64)
    assert D == 341
    assert gM == 31
    assert gL == 11


# ── Ratio E/(2φ)^n = 2φ/D exactly ────────────────────────────────────────────

def test_ratio_equals_two_phi_over_D():
    """E(n)/(2φ)^n = 2φ/D(n) to 1e-8 relative precision for n=36..80."""
    for n in range(36, 81):
        E = coupling_energy(n)
        D, _, _ = arithmetic_denominator(n)
        actual = E / TWO_PHI ** n
        predicted = TWO_PHI / D
        if predicted > 1e-15:
            rel = abs(actual - predicted) / predicted
            assert rel < 1e-8, \
                f"n={n}: ratio={actual}, 2φ/D={predicted}, rel={rel:.2e}"


# ── Mersenne × Lucas → (2φ)^{n+1} identity ───────────────────────────────────

def test_mersenne_lucas_product_identity():
    """M(k) × L(k) = (2φ)^k + (2ψ)^k - L(k) for k=1..50."""
    psi = (1 - math.sqrt(5)) / 2
    for k in range(1, 51):
        Mk = mersenne(k)
        Lk = lucas(k)
        lhs = Mk * Lk
        rhs = (2 * PHI) ** k + (2 * psi) ** k - Lk
        # rhs involves floats; lhs is exact integer
        assert abs(lhs - rhs) < max(1, abs(lhs) * 1e-10), \
            f"Identity failed at k={k}"
