"""
Tests for the theta module: Jacobi theta functions at the golden nome
and the Theta-Lucas dictionary.

Phase 88.0 | CSID: THETA-TEST-001
"""

import math
import pytest

from watkins_nn.constants import PHI, PHI_INV
from watkins_nn.kaleidoscope import lucas, fibonacci
from watkins_nn.theta import (
    golden_nome,
    golden_tau,
    theta2_golden,
    theta3_golden,
    theta4_golden,
    lucas_reciprocal_sum,
    fibonacci_reciprocal_sum,
    theta_lucas_identity,
    jacobi_triple_product,
    jacobi_triple_product_sum,
    eta_dedekind_golden,
    eta_quotient_golden,
    hecke_T_p,
    theta3_coefficients,
    theta3_squared_coefficients,
    delta_modular_golden,
    eisenstein_E4_golden,
    eisenstein_E6_golden,
    theta_lucas_report,
    lucas_reciprocal_tail_bound,
    theta_lucas_identity_precision,
    convergence_rate_comparison,
)


# =============================================================================
# GOLDEN NOME
# =============================================================================

class TestGoldenNome:
    def test_golden_nome_value(self):
        assert golden_nome() == pytest.approx(1.0 / PHI ** 2, abs=1e-15)

    def test_golden_nome_is_phi_inv_squared(self):
        assert golden_nome() == pytest.approx(PHI_INV ** 2, abs=1e-15)

    def test_golden_nome_less_than_one(self):
        assert 0 < golden_nome() < 1

    def test_golden_nome_numerical(self):
        assert golden_nome() == pytest.approx(0.3819660112501052, abs=1e-12)


# =============================================================================
# GOLDEN TAU
# =============================================================================

class TestGoldenTau:
    def test_tau_is_purely_imaginary(self):
        tau = golden_tau()
        assert tau.real == pytest.approx(0.0, abs=1e-15)
        assert tau.imag > 0

    def test_tau_gives_golden_nome(self):
        """e^{2*pi*i*tau} should equal the golden nome."""
        import cmath
        tau = golden_tau()
        q_recovered = cmath.exp(2 * math.pi * 1j * tau)
        assert q_recovered.real == pytest.approx(golden_nome(), abs=1e-12)
        assert abs(q_recovered.imag) < 1e-12


# =============================================================================
# THETA FUNCTIONS
# =============================================================================

class TestThetaFunctions:
    def test_theta3_positive(self):
        assert theta3_golden() > 1.0

    def test_theta2_positive(self):
        assert theta2_golden() > 0.0

    def test_theta4_positive(self):
        assert theta4_golden() > 0.0

    def test_theta3_value(self):
        assert theta3_golden() == pytest.approx(1.8068510462253213, abs=1e-10)

    def test_theta2_value(self):
        assert theta2_golden() == pytest.approx(1.8065967824482834, abs=1e-10)

    def test_theta4_value(self):
        assert theta4_golden() == pytest.approx(0.27829471999592426, abs=1e-10)

    def test_theta3_convergence(self):
        """More terms should not change the result significantly."""
        t50 = theta3_golden(50)
        t100 = theta3_golden(100)
        assert abs(t50 - t100) < 1e-14


# =============================================================================
# THETA-LUCAS IDENTITY
# =============================================================================

class TestThetaLucasIdentity:
    def test_identity_holds(self):
        """theta_3(q)^2 = 1 + 4 * sum 1/L(2n) at golden nome."""
        t3 = theta3_golden(100)
        lrs = lucas_reciprocal_sum(100)
        assert t3 ** 2 == pytest.approx(1.0 + 4.0 * lrs, abs=1e-12)

    def test_identity_dict(self):
        result = theta_lucas_identity(100)
        assert 'theta3_squared' in result
        assert 'lucas_side' in result
        assert 'difference' in result
        assert 'terms' in result
        assert result['terms'] == 100
        assert result['difference'] < 1e-12

    def test_lucas_reciprocal_sum_positive(self):
        assert lucas_reciprocal_sum(50) > 0

    def test_lucas_reciprocal_sum_convergence(self):
        s50 = lucas_reciprocal_sum(50)
        s100 = lucas_reciprocal_sum(100)
        assert abs(s50 - s100) < 1e-14

    def test_fibonacci_reciprocal_sum_positive(self):
        assert fibonacci_reciprocal_sum(50) > 0


# =============================================================================
# JACOBI IDENTITY: theta3^4 = theta2^4 + theta4^4
# =============================================================================

class TestJacobiIdentity:
    def test_jacobi_quartic_identity(self):
        t2 = theta2_golden(100)
        t3 = theta3_golden(100)
        t4 = theta4_golden(100)
        assert t3 ** 4 == pytest.approx(t2 ** 4 + t4 ** 4, abs=1e-10)

    def test_jacobi_quartic_values(self):
        t3 = theta3_golden(100)
        assert t3 ** 4 == pytest.approx(10.658335975885977, abs=1e-8)


# =============================================================================
# JACOBI TRIPLE PRODUCT
# =============================================================================

class TestJacobiTripleProduct:
    def test_triple_product_equals_series(self):
        """Product form should equal series form."""
        q = 0.3
        z = 1.0
        prod_val = jacobi_triple_product(q, z, 50)
        series_val = jacobi_triple_product_sum(q, z, 50)
        assert prod_val == pytest.approx(series_val, rel=1e-8)

    def test_triple_product_z1_is_theta3(self):
        """JTP with z=1 should give theta_3."""
        q = golden_nome()
        # theta3 = JTP(q^{1/2}, 1) when using the standard convention
        # Direct series: sum z^n q^{n^2} with z=1 is theta3
        series_val = jacobi_triple_product_sum(q, 1.0, 50)
        t3 = theta3_golden(50)
        assert series_val == pytest.approx(t3, abs=1e-10)

    def test_invalid_q_raises(self):
        with pytest.raises(ValueError):
            jacobi_triple_product(1.5, 1.0)

    def test_invalid_z_raises(self):
        with pytest.raises(ValueError):
            jacobi_triple_product(0.5, 0.0)


# =============================================================================
# DEDEKIND ETA
# =============================================================================

class TestDedekindEta:
    def test_eta_positive(self):
        assert eta_dedekind_golden() > 0

    def test_eta_value(self):
        assert eta_dedekind_golden(200) == pytest.approx(0.46251828769836895, abs=1e-10)

    def test_eta_quotient_self(self):
        """eta(tau)/eta(tau) = 1."""
        assert eta_quotient_golden(1, 1) == pytest.approx(1.0, abs=1e-10)

    def test_delta_modular_positive(self):
        assert delta_modular_golden() > 0

    def test_delta_equals_eta_24(self):
        eta = eta_dedekind_golden(100)
        delta = delta_modular_golden(100)
        assert delta == pytest.approx(eta ** 24, rel=1e-6)


# =============================================================================
# Q-EXPANSION COEFFICIENTS
# =============================================================================

class TestCoefficients:
    def test_theta3_coefficients_constant(self):
        c = theta3_coefficients(10)
        assert c[0] == 1

    def test_theta3_coefficients_squares(self):
        c = theta3_coefficients(10)
        assert c[1] == 2  # q^1 = q^{1^2}
        assert c[4] == 2  # q^4 = q^{2^2}
        assert c[9] == 2  # q^9 = q^{3^2}

    def test_theta3_coefficients_nonsquares(self):
        c = theta3_coefficients(10)
        assert c[2] == 0
        assert c[3] == 0
        assert c[5] == 0

    def test_theta3_squared_coefficients_r2(self):
        """Coefficient of q^n in theta3^2 is r_2(n), representations as sum of 2 squares."""
        c = theta3_squared_coefficients(10)
        assert c[0] == 1  # 0 = 0^2 + 0^2
        assert c[1] == 4  # 1 = (+-1)^2 + 0^2 or 0^2 + (+-1)^2
        assert c[2] == 4  # 2 = (+-1)^2 + (+-1)^2
        assert c[5] == 8  # 5 = (+-1)^2 + (+-2)^2 or (+-2)^2 + (+-1)^2


# =============================================================================
# HECKE OPERATORS
# =============================================================================

class TestHeckeOperators:
    def test_hecke_T2_on_theta3(self):
        """T_2 should produce valid coefficients."""
        c = theta3_coefficients(20)
        c_float = [float(x) for x in c]
        result = hecke_T_p(2, c_float)
        assert len(result) > 0
        assert isinstance(result[0], float)

    def test_hecke_T2_coefficient_0(self):
        """(T_2 f)_0 = a_0 + a_0 = 2*a_0 for theta3."""
        c = [float(x) for x in theta3_coefficients(20)]
        result = hecke_T_p(2, c)
        # n=0: a_{2*0} + a_{0/2} = a_0 + a_0 = 2
        assert result[0] == pytest.approx(2.0)

    def test_hecke_T2_coefficient_1(self):
        """(T_2 f)_1 = a_2 for theta3 (since 2 does not divide 1)."""
        c = [float(x) for x in theta3_coefficients(20)]
        result = hecke_T_p(2, c)
        # n=1: a_{2} + 0 = 0
        assert result[1] == pytest.approx(0.0)

    def test_hecke_invalid_p(self):
        with pytest.raises(ValueError):
            hecke_T_p(1, [1.0, 0.0, 1.0])


# =============================================================================
# EISENSTEIN SERIES
# =============================================================================

class TestEisenstein:
    def test_e4_greater_than_1(self):
        assert eisenstein_E4_golden(50) > 1.0

    def test_e6_less_than_1(self):
        # E6 starts at 1 but subtracts, so for large enough q it should be < 1
        assert eisenstein_E6_golden(50) < 1.0

    def test_e4_e6_discriminant_relation(self):
        """Delta = (E4^3 - E6^2) / 1728."""
        e4 = eisenstein_E4_golden(80)
        e6 = eisenstein_E6_golden(80)
        delta_from_eis = (e4 ** 3 - e6 ** 2) / 1728.0
        delta_direct = delta_modular_golden(80)
        assert delta_from_eis == pytest.approx(delta_direct, rel=0.05)


# =============================================================================
# FULL REPORT
# =============================================================================

class TestThetaLucasReport:
    def test_report_keys(self):
        report = theta_lucas_report(20)
        expected_keys = {
            'golden_nome', 'golden_tau', 'theta2', 'theta3', 'theta4',
            'theta3_squared', 'lucas_sum', 'lucas_identity_rhs',
            'lucas_identity_error', 'jacobi_identity_lhs',
            'jacobi_identity_rhs', 'jacobi_identity_error',
            'eta_dedekind', 'theta2_over_theta3', 'theta4_over_theta3',
        }
        assert set(report.keys()) == expected_keys

    def test_report_identities_hold(self):
        report = theta_lucas_report(50)
        assert report['lucas_identity_error'] < 1e-12
        assert report['jacobi_identity_error'] < 1e-10


# =============================================================================
# TAIL BOUNDS FOR LUCAS RECIPROCAL SUM
# =============================================================================

class TestLucasReciprocalTailBound:
    def test_bounds_are_positive(self):
        lower, upper = lucas_reciprocal_tail_bound(10)
        assert lower > 0
        assert upper > 0

    def test_lower_leq_upper(self):
        for N in [5, 10, 20, 50, 100]:
            lower, upper = lucas_reciprocal_tail_bound(N)
            assert lower <= upper, f"lower > upper at N={N}"

    def test_bounds_bracket_actual_tail(self):
        """Verify lower <= actual tail <= upper for N where float64 resolves the tail.

        For large N the tail is below float64 precision relative to the full sum,
        so we only test small N where subtraction is still meaningful.
        We allow a small relative tolerance for float64 rounding.
        """
        # Compute a high-accuracy full sum as reference
        full_sum = lucas_reciprocal_sum(500)
        for N in [5, 10, 15]:
            partial = lucas_reciprocal_sum(N)
            actual_tail = full_sum - partial
            lower, upper = lucas_reciprocal_tail_bound(N)
            # Allow relative tolerance for float64 rounding
            tol = max(abs(upper), abs(actual_tail)) * 1e-6
            assert lower <= actual_tail + tol, (
                f"N={N}: lower={lower} > actual={actual_tail}"
            )
            assert actual_tail <= upper + tol, (
                f"N={N}: actual={actual_tail} > upper={upper}"
            )

    def test_bounds_shrink_with_N(self):
        """Bounds should decrease as N increases."""
        _, upper_10 = lucas_reciprocal_tail_bound(10)
        _, upper_20 = lucas_reciprocal_tail_bound(20)
        _, upper_50 = lucas_reciprocal_tail_bound(50)
        assert upper_20 < upper_10
        assert upper_50 < upper_20

    def test_tail_bound_N100_very_small(self):
        """At N=100, the tail should be astronomically small."""
        _, upper = lucas_reciprocal_tail_bound(100)
        assert upper < 1e-40

    def test_bounds_tight(self):
        """The gap (upper - lower) should be much smaller than the upper bound."""
        lower, upper = lucas_reciprocal_tail_bound(10)
        gap = upper - lower
        # The correction term (gap) should be much smaller than the leading term
        assert gap < 0.01 * upper


# =============================================================================
# THETA-LUCAS IDENTITY PRECISION
# =============================================================================

class TestThetaLucasIdentityPrecision:
    def test_returns_tuple(self):
        val, err = theta_lucas_identity_precision(50)
        assert isinstance(val, float)
        assert isinstance(err, float)

    def test_identity_value_positive(self):
        """With finite terms, theta3^2 > lucas_side (missing positive tail)."""
        val, _ = theta_lucas_identity_precision(20)
        # The identity value = theta3^2 - lucas_side
        # theta3 converges much faster, so its truncation is negligible.
        # The Lucas side is missing positive tail terms, so lucas_side < truth.
        # Therefore theta3^2 ≈ truth > lucas_side, giving val > 0.
        assert val > 0

    def test_error_bound_positive(self):
        _, err = theta_lucas_identity_precision(50)
        assert err > 0

    def test_identity_value_within_error_bound(self):
        """The computed difference should be bounded by the error estimate."""
        for terms in [10, 20, 50, 100]:
            val, err = theta_lucas_identity_precision(terms)
            # For high term counts, float64 arithmetic floors at ~1e-15;
            # the predicted tail bound can be far below this floor.
            effective_bound = max(err, 1e-14) * 1.01
            assert abs(val) <= effective_bound, (
                f"terms={terms}: |val|={abs(val)} > bound={effective_bound}"
            )

    def test_error_shrinks_with_terms(self):
        _, err_10 = theta_lucas_identity_precision(10)
        _, err_50 = theta_lucas_identity_precision(50)
        _, err_100 = theta_lucas_identity_precision(100)
        assert err_50 < err_10
        assert err_100 < err_50

    def test_error_matches_tail_prediction(self):
        """The identity error should approximately equal 4 * tail."""
        for N in [10, 20, 50]:
            val, _ = theta_lucas_identity_precision(N)
            _, tail_upper = lucas_reciprocal_tail_bound(N)
            # val ≈ 4 * actual_tail, and actual_tail ≤ tail_upper.
            # Floor at machine epsilon: float64 cannot resolve below ~1e-15.
            bound = max(4.0 * tail_upper, 1e-14) * 1.01
            assert abs(val) <= bound

    def test_high_precision_at_200_terms(self):
        """200 terms should give essentially machine-epsilon accuracy."""
        val, err = theta_lucas_identity_precision(200)
        assert abs(val) < 1e-14
        assert err < 1e-14


# =============================================================================
# CONVERGENCE RATE COMPARISON
# =============================================================================

class TestConvergenceRateComparison:
    def test_returns_dict(self):
        result = convergence_rate_comparison(10)
        assert isinstance(result, dict)

    def test_required_keys(self):
        result = convergence_rate_comparison(10)
        expected = {
            'theta3_term_N', 'lucas_term_N',
            'theta3_log_rate', 'lucas_log_rate',
            'theta3_terms_for_eps', 'lucas_terms_for_eps',
            'bottleneck', 'terms_for_1e15',
        }
        assert set(result.keys()) == expected

    def test_theta3_converges_faster(self):
        """θ₃ term at n=10 should be much smaller than Lucas term."""
        result = convergence_rate_comparison(10)
        assert result['theta3_term_N'] < result['lucas_term_N']

    def test_theta3_needs_fewer_terms(self):
        result = convergence_rate_comparison(10)
        assert result['theta3_terms_for_eps'] < result['lucas_terms_for_eps']

    def test_bottleneck_is_lucas(self):
        """Lucas sum should be the bottleneck for precision."""
        result = convergence_rate_comparison(10)
        assert result['bottleneck'] == 'lucas'

    def test_lucas_terms_for_eps_reasonable(self):
        """Should need roughly 36-38 terms for machine epsilon."""
        result = convergence_rate_comparison(10)
        assert 30 <= result['lucas_terms_for_eps'] <= 45

    def test_theta3_terms_for_eps_small(self):
        """θ₃ needs very few terms (< 10) for machine epsilon."""
        result = convergence_rate_comparison(10)
        assert result['theta3_terms_for_eps'] <= 10

    def test_terms_for_1e15_reasonable(self):
        result = convergence_rate_comparison(10)
        assert 30 <= result['terms_for_1e15'] <= 40

    def test_log_rates_positive(self):
        result = convergence_rate_comparison(10)
        assert result['theta3_log_rate'] > 0
        assert result['lucas_log_rate'] > 0

    def test_theta3_log_rate_much_higher(self):
        """At n=10, θ₃ has gained many more digits than Lucas."""
        result = convergence_rate_comparison(10)
        # theta3 at n=10: ~100 * log10(1/q) ≈ 42 digits
        # lucas at n=10: ~20 * log10(phi) ≈ 4 digits
        assert result['theta3_log_rate'] > 5 * result['lucas_log_rate']
