"""
Tests for the Hecke eigenform property of theta_3^2 and the Lucas form.

FINDINGS:
=========
1. theta_3^2 IS a Hecke eigenform for all odd primes p not dividing the level 4.
   Eigenvalue: lambda_p = 1 + chi_{-4}(p)
     = 2  if p ≡ 1 (mod 4)
     = 0  if p ≡ 3 (mod 4)
   This identifies theta_3^2 as the weight-1 Eisenstein series E_1(chi_0, chi_{-4})
   on Gamma_0(4).

2. For p=2 (bad prime, 2 | level 4): theta_3^2 is an eigenform of the Atkin
   operator U_2 (defined by (U_2 f)_n = a_{2n}) with eigenvalue 1.
   It is NOT an eigenform of the full T_2 operator because T_2 includes the
   back-shift term a_{n/2} which mixes coefficients from different chi_{-4} classes.

3. The Lucas-specific form f(q) = sum_{n>=1} q^n / L(2n) is NOT an eigenform.
   Its coefficients 1/L(2n) are not multiplicative (violates necessary condition).

Phase 88.0 | CSID: HECKE-EIGENFORM-001
"""

import math
import pytest

from watkins_nn.theta import (
    theta3_squared_coefficients,
    theta3_coefficients,
    hecke_T_p,
    golden_nome,
    is_hecke_eigenform,
    chi_minus4,
)
from watkins_nn.constants import PHI
from watkins_nn.kaleidoscope import lucas


# =============================================================================
# THETA_3^2 EIGENFORM TESTS
# =============================================================================

class TestTheta3SquaredEigenform:
    """theta_3^2 = sum r_2(n) q^n is a Hecke eigenform with lambda_p = 1 + chi_{-4}(p)."""

    @pytest.fixture
    def coeffs(self):
        return theta3_squared_coefficients(600)

    def test_eigenvalue_formula_odd_primes(self, coeffs):
        """Verify lambda_p = 1 + chi_{-4}(p) for many odd primes."""
        a = [float(x) for x in coeffs]
        primes = [3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47]

        for p in primes:
            expected = 1 + chi_minus4(p)
            Tp_a = hecke_T_p(p, a)
            # Check at all nonzero coefficients (n >= 1)
            for n in range(1, len(Tp_a)):
                if abs(a[n]) > 0.5:
                    actual = Tp_a[n] / a[n]
                    assert abs(actual - expected) < 1e-10, \
                        f"p={p}, n={n}: expected {expected}, got {actual}"

    def test_eigenvalue_p1_mod4_is_2(self, coeffs):
        """Primes p ≡ 1 mod 4 have eigenvalue 2."""
        a = [float(x) for x in coeffs]
        for p in [5, 13, 17, 29, 37, 41]:
            Tp_a = hecke_T_p(p, a)
            for n in range(1, min(len(Tp_a), 50)):
                if abs(a[n]) > 0.5:
                    assert abs(Tp_a[n] - 2.0 * a[n]) < 1e-10, \
                        f"T_{p} not doubling at n={n}"

    def test_eigenvalue_p3_mod4_is_0(self, coeffs):
        """Primes p ≡ 3 mod 4 annihilate the form (eigenvalue 0)."""
        a = [float(x) for x in coeffs]
        for p in [3, 7, 11, 19, 23, 31, 43, 47]:
            Tp_a = hecke_T_p(p, a)
            for n in range(1, min(len(Tp_a), 50)):
                if abs(a[n]) > 0.5:
                    assert abs(Tp_a[n]) < 1e-10, \
                        f"T_{p} should annihilate at n={n}, got {Tp_a[n]}"

    def test_U2_eigenvalue_is_1(self, coeffs):
        """Atkin operator U_2: (U_2 f)_n = a_{2n} has eigenvalue 1."""
        a = [float(x) for x in coeffs]
        for n in range(1, 200):
            if abs(a[n]) > 0.5:
                assert abs(a[2 * n] - a[n]) < 1e-10, \
                    f"U_2 eigenvalue not 1 at n={n}: a[{2*n}]={a[2*n]}, a[{n}]={a[n]}"

    def test_T2_not_scalar_eigenform(self, coeffs):
        """Full T_2 is NOT a scalar eigenform (ratios are not constant)."""
        a = [float(x) for x in coeffs]
        Tp_a = hecke_T_p(2, a)
        ratios = set()
        for n in range(1, min(len(Tp_a), 50)):
            if abs(a[n]) > 0.5:
                ratios.add(round(Tp_a[n] / a[n], 8))
        # Ratios should include both 1.0 and 2.0
        assert len(ratios) > 1, "T_2 ratios should not all be equal"

    def test_is_hecke_eigenform_function(self, coeffs):
        """Test the is_hecke_eigenform convenience function."""
        result = is_hecke_eigenform(coeffs)
        assert result['is_eigenform'] is True
        assert result['eigenvalue_formula'] == '1 + chi_{-4}(p)'
        # Eigenvalues for specific primes
        assert result['eigenvalues'][5] == pytest.approx(2.0, abs=1e-10)
        assert result['eigenvalues'][3] == pytest.approx(0.0, abs=1e-10)
        assert result['eigenvalues'][13] == pytest.approx(2.0, abs=1e-10)
        assert result['eigenvalues'][7] == pytest.approx(0.0, abs=1e-10)


# =============================================================================
# CHI_{-4} CHARACTER TESTS
# =============================================================================

class TestChiMinus4:
    """Test the Dirichlet character chi_{-4} = Kronecker symbol (-4/.)."""

    def test_chi_at_primes(self):
        assert chi_minus4(2) == 0
        assert chi_minus4(3) == -1
        assert chi_minus4(5) == 1
        assert chi_minus4(7) == -1
        assert chi_minus4(11) == -1
        assert chi_minus4(13) == 1
        assert chi_minus4(17) == 1
        assert chi_minus4(19) == -1

    def test_chi_at_composites(self):
        """chi_{-4} is completely multiplicative."""
        assert chi_minus4(1) == 1
        assert chi_minus4(4) == 0    # 4 = 2*2, even
        assert chi_minus4(9) == 1    # 9 ≡ 1 mod 4
        assert chi_minus4(15) == -1  # 15 ≡ 3 mod 4, also chi(3)*chi(5) = (-1)(1) = -1
        assert chi_minus4(21) == 1   # 21 ≡ 1 mod 4, also chi(3)*chi(7) = (-1)(-1) = 1
        assert chi_minus4(25) == 1   # 25 ≡ 1 mod 4

    def test_chi_periodicity(self):
        """chi_{-4} has period 4: chi(n) depends only on n mod 4."""
        for n in range(1, 100):
            if n % 2 == 0:
                assert chi_minus4(n) == 0
            elif n % 4 == 1:
                assert chi_minus4(n) == 1
            else:
                assert chi_minus4(n) == -1


# =============================================================================
# LUCAS FORM: NOT AN EIGENFORM
# =============================================================================

class TestLucasFormNotEigenform:
    """f(q) = sum_{n>=1} q^n / L(2n) is NOT a Hecke eigenform."""

    def _lucas_coefficients(self, max_n):
        coeffs = [0.0] * (max_n + 1)
        for n in range(1, max_n + 1):
            coeffs[n] = 1.0 / lucas(2 * n)
        return coeffs

    def test_not_multiplicative(self):
        """Eigenform coefficients must be multiplicative; Lucas coeffs are not."""
        c = self._lucas_coefficients(100)
        # a(2)*a(3) should equal a(6) for multiplicativity
        product = c[2] * c[3]
        actual = c[6]
        assert abs(product - actual) / abs(actual) > 0.1, \
            "Lucas coefficients should NOT be multiplicative"

    def test_hecke_ratios_not_constant(self):
        """T_p ratios are wildly non-constant for the Lucas form."""
        c = self._lucas_coefficients(200)
        for p in [3, 5, 7]:
            Tp_c = hecke_T_p(p, c)
            ratios = []
            for n in range(1, min(len(Tp_c), 30)):
                if abs(c[n]) > 1e-50:
                    ratios.append(Tp_c[n] / c[n])
            if len(ratios) >= 2:
                # Ratios should vary wildly
                spread = max(ratios) - min(ratios)
                assert spread > 0.01, \
                    f"T_{p} ratios too constant for Lucas form: spread={spread}"


# =============================================================================
# R_2(N) FORMULA VERIFICATION
# =============================================================================

class TestR2Formula:
    """Verify r_2(n) = 4 * sum_{d|n} chi_{-4}(d) (the divisor sum formula)."""

    def test_r2_divisor_sum_formula(self):
        """r_2(n) = 4 * sum_{d|n} chi_{-4}(d) for n >= 1."""
        coeffs = theta3_squared_coefficients(100)
        for n in range(1, 101):
            divisor_sum = sum(chi_minus4(d) for d in range(1, n + 1) if n % d == 0)
            expected = 4 * divisor_sum
            assert coeffs[n] == expected, \
                f"r_2({n}): got {coeffs[n]}, expected 4*{divisor_sum}={expected}"
