"""
Jacobi theta functions at the golden nome and the Theta-Lucas dictionary.

The golden nome q = 1/phi^2 connects Jacobi theta functions to Lucas numbers
via the identity:

    theta_3(q)^2 = 1 + 4 * sum_{n=1}^inf 1/L(2n)

where L(n) is the n-th Lucas number.

This module implements:
- Classical Jacobi theta functions theta_2, theta_3, theta_4 evaluated at q = 1/phi^2
- The Theta-Lucas dictionary relating modular forms to Fibonacci/Lucas arithmetic
- Dedekind eta function at the golden nome
- Jacobi triple product
- Hecke operators on q-expansions

Phase 88.0 | CSID: THETA-001
"""

import math
from typing import Dict, List

from watkins_nn.constants import PHI, PHI_INV
from watkins_nn.kaleidoscope import lucas, fibonacci


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

def golden_nome() -> float:
    """Return the golden nome q = 1/phi^2.

    This is the unique nome where Jacobi theta functions connect
    to Lucas number reciprocal sums.
    """
    return 1.0 / (PHI ** 2)


# =============================================================================
# JACOBI THETA FUNCTIONS AT THE GOLDEN NOME
# =============================================================================

def theta2_golden(terms: int = 100) -> float:
    """Jacobi theta_2 at the golden nome.

    theta_2(q) = 2 * q^{1/4} * sum_{n=0}^{N} q^{n(n+1)}

    Parameters
    ----------
    terms : int
        Number of terms in the series.

    Returns
    -------
    float
        Value of theta_2 at q = 1/phi^2.
    """
    q = golden_nome()
    return 2.0 * q ** 0.25 * sum(q ** (n * (n + 1)) for n in range(terms))


def theta3_golden(terms: int = 100) -> float:
    """Jacobi theta_3 at the golden nome.

    theta_3(q) = 1 + 2 * sum_{n=1}^{N} q^{n^2}

    Parameters
    ----------
    terms : int
        Number of terms in the series.

    Returns
    -------
    float
        Value of theta_3 at q = 1/phi^2.
    """
    q = golden_nome()
    return 1.0 + 2.0 * sum(q ** (n * n) for n in range(1, terms + 1))


def theta4_golden(terms: int = 100) -> float:
    """Jacobi theta_4 at the golden nome.

    theta_4(q) = 1 + 2 * sum_{n=1}^{N} (-1)^n * q^{n^2}

    Parameters
    ----------
    terms : int
        Number of terms in the series.

    Returns
    -------
    float
        Value of theta_4 at q = 1/phi^2.
    """
    q = golden_nome()
    return 1.0 + 2.0 * sum((-1) ** n * q ** (n * n) for n in range(1, terms + 1))


# =============================================================================
# THETA-LUCAS DICTIONARY
# =============================================================================

def lucas_reciprocal_sum(terms: int = 100) -> float:
    """Compute sum_{n=1}^{N} 1/L(2n).

    This sum connects to theta_3 via the identity:
        theta_3(q)^2 = 1 + 4 * sum_{n=1}^inf 1/L(2n)
    at the golden nome q = 1/phi^2.

    Parameters
    ----------
    terms : int
        Number of terms.

    Returns
    -------
    float
        The partial sum.
    """
    return sum(1.0 / lucas(2 * n) for n in range(1, terms + 1))


def theta_lucas_identity(terms: int = 100) -> Dict[str, float]:
    """Verify the Theta-Lucas identity.

    theta_3(q)^2 = 1 + 4 * sum_{n=1}^inf 1/L(2n)

    Returns
    -------
    dict
        Keys: 'theta3_squared', 'lucas_side', 'difference', 'terms'
    """
    t3 = theta3_golden(terms)
    lrs = lucas_reciprocal_sum(terms)
    theta3_sq = t3 ** 2
    lucas_side = 1.0 + 4.0 * lrs
    return {
        'theta3_squared': theta3_sq,
        'lucas_side': lucas_side,
        'difference': abs(theta3_sq - lucas_side),
        'terms': terms,
    }


def fibonacci_reciprocal_sum(terms: int = 100) -> float:
    """Compute sum_{n=1}^{N} 1/F(2n).

    Companion to the Lucas reciprocal sum in the Theta-Lucas dictionary.

    Parameters
    ----------
    terms : int
        Number of terms.

    Returns
    -------
    float
        The partial sum.
    """
    return sum(1.0 / fibonacci(2 * n) for n in range(1, terms + 1))


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

def jacobi_triple_product(q: float, z: float, terms: int = 100) -> float:
    """Jacobi triple product.

    JTP(q, z) = prod_{n=1}^{N} (1 - q^{2n})(1 + z*q^{2n-1})(1 + z^{-1}*q^{2n-1})

    This equals sum_{n=-N}^{N} z^n * q^{n^2} (the theta function series).

    Parameters
    ----------
    q : float
        The nome (|q| < 1).
    z : float
        The argument (z != 0).
    terms : int
        Number of product terms.

    Returns
    -------
    float
        Value of the triple product.
    """
    if abs(q) >= 1:
        raise ValueError("Nome q must satisfy |q| < 1")
    if z == 0:
        raise ValueError("z must be nonzero")
    product = 1.0
    z_inv = 1.0 / z
    for n in range(1, terms + 1):
        q2n = q ** (2 * n)
        q2nm1 = q ** (2 * n - 1)
        product *= (1 - q2n) * (1 + z * q2nm1) * (1 + z_inv * q2nm1)
    return product


def jacobi_triple_product_sum(q: float, z: float, terms: int = 100) -> float:
    """Jacobi triple product via the series representation.

    sum_{n=-N}^{N} z^n * q^{n^2}

    Parameters
    ----------
    q : float
        The nome.
    z : float
        The argument.
    terms : int
        Number of terms on each side of n=0.

    Returns
    -------
    float
        Value of the series.
    """
    return sum(z ** n * q ** (n * n) for n in range(-terms, terms + 1))


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

def golden_tau() -> complex:
    """Return tau such that q = e^{2*pi*i*tau} equals the golden nome.

    For q = 1/phi^2 (real, positive, < 1), tau is purely imaginary:
        tau = i * ln(phi^2) / (2*pi) = i * ln(phi^2) / (2*pi)

    Returns
    -------
    complex
        The value of tau (purely imaginary, positive imaginary part).
    """
    q = golden_nome()
    return 1j * (-math.log(q)) / (2 * math.pi)


def eta_dedekind_golden(terms: int = 200) -> float:
    """Dedekind eta function at the golden nome.

    eta(tau) = q^{1/24} * prod_{n=1}^{N} (1 - q^n)

    where q = e^{2*pi*i*tau} = 1/phi^2.

    Parameters
    ----------
    terms : int
        Number of product terms.

    Returns
    -------
    float
        Value of eta(tau) at the golden nome.
    """
    q = golden_nome()
    return q ** (1.0 / 24.0) * math.prod(1 - q ** n for n in range(1, terms + 1))


def eta_quotient_golden(a: int = 1, b: int = 1, terms: int = 200) -> float:
    """Compute eta(a*tau)/eta(b*tau) at the golden nome.

    Eta quotients are building blocks for modular forms of various levels.

    Parameters
    ----------
    a, b : int
        Multipliers for tau.
    terms : int
        Number of product terms.

    Returns
    -------
    float
        Value of eta(a*tau)/eta(b*tau).
    """
    q = golden_nome()
    qa = q ** a
    qb = q ** b
    prefactor = qa ** (1.0 / 24.0) / qb ** (1.0 / 24.0)
    num = math.prod(1 - qa ** n for n in range(1, terms + 1))
    den = math.prod(1 - qb ** n for n in range(1, terms + 1))
    return prefactor * num / den


# =============================================================================
# THETA SERIES COEFFICIENTS (q-EXPANSION)
# =============================================================================

def theta3_coefficients(max_n: int) -> List[int]:
    """Return the q-expansion coefficients of theta_3.

    theta_3(q) = sum_{n=-inf}^{inf} q^{n^2} = 1 + 2*q + 2*q^4 + 2*q^9 + ...

    The coefficient of q^k is the number of representations of k as n^2
    with sign, i.e. r_1(k) counting +-n.

    Parameters
    ----------
    max_n : int
        Maximum exponent in the q-expansion.

    Returns
    -------
    list of int
        Coefficients [a_0, a_1, ..., a_{max_n}].
    """
    coeffs = [0] * (max_n + 1)
    coeffs[0] = 1
    n = 1
    while n * n <= max_n:
        coeffs[n * n] += 2
        n += 1
    return coeffs


def theta3_squared_coefficients(max_n: int) -> List[int]:
    """Return q-expansion coefficients of theta_3^2.

    theta_3^2 = (sum q^{n^2})^2 = sum r_2(k) q^k

    where r_2(k) is the number of representations of k as a sum of two squares.

    Parameters
    ----------
    max_n : int
        Maximum exponent.

    Returns
    -------
    list of int
        Coefficients [a_0, a_1, ..., a_{max_n}].
    """
    c = theta3_coefficients(max_n)
    result = [0] * (max_n + 1)
    for i in range(max_n + 1):
        if c[i] == 0:
            continue
        for j in range(max_n + 1 - i):
            if c[j] == 0:
                continue
            result[i + j] += c[i] * c[j]
    return result


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

def hecke_T_p(p: int, coeffs: List[float]) -> List[float]:
    """Apply Hecke operator T_p to a q-expansion.

    For a modular form f = sum a_n q^n of weight k, the Hecke operator acts as:
        (T_p f)_n = a_{pn} + p^{k-1} * a_{n/p}

    where a_{n/p} = 0 if p does not divide n.

    We use weight k=1 (theta functions), so p^{k-1} = 1.

    Parameters
    ----------
    p : int
        A prime number.
    coeffs : list of float
        The q-expansion coefficients [a_0, a_1, ...].

    Returns
    -------
    list of float
        The coefficients of T_p(f).

    Raises
    ------
    ValueError
        If p < 2.
    """
    if p < 2:
        raise ValueError("p must be >= 2")
    N = len(coeffs)
    result_len = N // p  # T_p reduces the number of usable coefficients
    if result_len == 0:
        return []
    result = [0.0] * result_len
    for n in range(result_len):
        # a_{pn} term
        pn = p * n
        if pn < N:
            result[n] += coeffs[pn]
        # p^{k-1} * a_{n/p} term (k=1 so coefficient is 1)
        if n % p == 0:
            idx = n // p
            if idx < N:
                result[n] += coeffs[idx]
    return result


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

def chi_minus4(n: int) -> int:
    """Dirichlet character chi_{-4} = Kronecker symbol (-4/n).

    This is the unique primitive character of conductor 4:
        chi_{-4}(n) =  0  if n is even
        chi_{-4}(n) = +1  if n ≡ 1 (mod 4)
        chi_{-4}(n) = -1  if n ≡ 3 (mod 4)

    Completely multiplicative: chi(ab) = chi(a) * chi(b).

    Parameters
    ----------
    n : int
        Positive integer.

    Returns
    -------
    int
        Value of chi_{-4}(n) in {-1, 0, 1}.
    """
    if n <= 0:
        raise ValueError("n must be positive")
    if n % 2 == 0:
        return 0
    return 1 if n % 4 == 1 else -1


# =============================================================================
# HECKE EIGENFORM DETECTION
# =============================================================================

def is_hecke_eigenform(coeffs: List[int], primes: List[int] = None,
                       max_check: int = 50) -> Dict:
    """Test whether theta_3^2 coefficients form a Hecke eigenform.

    theta_3^2 = sum r_2(n) q^n is the weight-1 Eisenstein series E_1(chi_0, chi_{-4})
    on Gamma_0(4). Its Hecke eigenvalues are:

        lambda_p = 1 + chi_{-4}(p)

    This function verifies this identity numerically for the given coefficients.

    Parameters
    ----------
    coeffs : list of int
        The q-expansion coefficients [a_0, a_1, ...] (should be r_2(n) values).
    primes : list of int, optional
        Primes to test. Default: [3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47].
    max_check : int
        Maximum index n to check for each prime.

    Returns
    -------
    dict
        Keys:
        - 'is_eigenform': bool — True if eigenform property holds
        - 'eigenvalue_formula': str — the formula for eigenvalues
        - 'eigenvalues': dict — {p: lambda_p} for each tested prime
        - 'max_deviation': float — largest deviation from predicted eigenvalue
        - 'form_identification': str — identification as Eisenstein series
    """
    if primes is None:
        primes = [3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47]

    a = [float(x) for x in coeffs]
    eigenvalues = {}
    max_deviation = 0.0
    is_eigen = True

    for p in primes:
        expected_lam = 1 + chi_minus4(p)
        Tp_a = hecke_T_p(p, a)

        # Compute actual eigenvalue from nonzero coefficients (n >= 1)
        ratios = []
        for n in range(1, min(len(Tp_a), max_check)):
            if abs(a[n]) > 0.5:
                ratio = Tp_a[n] / a[n]
                ratios.append(ratio)
                dev = abs(ratio - expected_lam)
                if dev > max_deviation:
                    max_deviation = dev
                if dev > 1e-8:
                    is_eigen = False

        if ratios:
            eigenvalues[p] = sum(ratios) / len(ratios)
        else:
            eigenvalues[p] = float('nan')

    return {
        'is_eigenform': is_eigen,
        'eigenvalue_formula': '1 + chi_{-4}(p)',
        'eigenvalues': eigenvalues,
        'max_deviation': max_deviation,
        'form_identification': (
            'Weight-1 Eisenstein series E_1(chi_0, chi_{-4}) on Gamma_0(4). '
            'Equivalently, theta_3(q)^2 = 1 + 4*sum_{d|n} chi_{-4}(d).'
        ),
    }


# =============================================================================
# MODULAR DISCRIMINANT
# =============================================================================

def delta_modular_golden(terms: int = 100) -> float:
    """Ramanujan's modular discriminant at the golden nome.

    Delta(tau) = eta(tau)^24 = q * prod_{n=1}^inf (1-q^n)^24

    Parameters
    ----------
    terms : int
        Number of product terms.

    Returns
    -------
    float
        Value of Delta at the golden nome.
    """
    q = golden_nome()
    prod_val = math.prod((1 - q ** n) for n in range(1, terms + 1))
    return q * prod_val ** 24


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

def eisenstein_E4_golden(terms: int = 100) -> float:
    """Eisenstein series E_4 at the golden nome via q-expansion.

    E_4(tau) = 1 + 240 * sum_{n=1}^N sigma_3(n) * q^n

    where sigma_3(n) = sum of cubes of divisors of n.

    Parameters
    ----------
    terms : int
        Number of terms.

    Returns
    -------
    float
        Value of E_4 at q = 1/phi^2.
    """
    q = golden_nome()
    total = 0.0
    for n in range(1, terms + 1):
        s3 = sum(d ** 3 for d in range(1, n + 1) if n % d == 0)
        total += s3 * q ** n
    return 1.0 + 240.0 * total


def eisenstein_E6_golden(terms: int = 100) -> float:
    """Eisenstein series E_6 at the golden nome via q-expansion.

    E_6(tau) = 1 - 504 * sum_{n=1}^N sigma_5(n) * q^n

    Parameters
    ----------
    terms : int
        Number of terms.

    Returns
    -------
    float
        Value of E_6 at q = 1/phi^2.
    """
    q = golden_nome()
    total = 0.0
    for n in range(1, terms + 1):
        s5 = sum(d ** 5 for d in range(1, n + 1) if n % d == 0)
        total += s5 * q ** n
    return 1.0 - 504.0 * total


# =============================================================================
# CONVENIENCE: FULL THETA-LUCAS REPORT
# =============================================================================

def theta_lucas_report(terms: int = 100) -> Dict[str, float]:
    """Complete Theta-Lucas dictionary evaluation.

    Returns all theta values, Lucas sums, identities, and derived quantities.

    Parameters
    ----------
    terms : int
        Number of terms for all series.

    Returns
    -------
    dict
        Comprehensive dictionary of all computed values.
    """
    t2 = theta2_golden(terms)
    t3 = theta3_golden(terms)
    t4 = theta4_golden(terms)
    lrs = lucas_reciprocal_sum(terms)
    eta = eta_dedekind_golden(2 * terms)

    return {
        'golden_nome': golden_nome(),
        'golden_tau': golden_tau().imag,
        'theta2': t2,
        'theta3': t3,
        'theta4': t4,
        'theta3_squared': t3 ** 2,
        'lucas_sum': lrs,
        'lucas_identity_rhs': 1.0 + 4.0 * lrs,
        'lucas_identity_error': abs(t3 ** 2 - (1.0 + 4.0 * lrs)),
        'jacobi_identity_lhs': t3 ** 4,
        'jacobi_identity_rhs': t2 ** 4 + t4 ** 4,
        'jacobi_identity_error': abs(t3 ** 4 - (t2 ** 4 + t4 ** 4)),
        'eta_dedekind': eta,
        'theta2_over_theta3': t2 / t3,
        'theta4_over_theta3': t4 / t3,
    }


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

def lucas_reciprocal_tail_bound(N: int) -> tuple:
    """Rigorous upper and lower bounds for Σ_{n>N} 1/L(2n).

    Derivation:
        L(2n) = φ^{2n} + ψ^{2n} = φ^{2n}(1 + φ^{-4n})
        1/L(2n) = φ^{-2n} / (1 + φ^{-4n})
                = φ^{-2n} · Σ_{k=0}^∞ (-1)^k φ^{-4kn}    (geometric series, |φ^{-4n}| < 1)
                = φ^{-2n} - φ^{-6n} + φ^{-10n} - ...

    Summing over n > N:
        R_N = Σ_{n>N} φ^{-2n}  - Σ_{n>N} φ^{-6n}  + Σ_{n>N} φ^{-10n} - ...
            = φ^{-2(N+1)}/(1 - φ^{-2}) - φ^{-6(N+1)}/(1 - φ^{-6}) + ...

    The alternating series gives:
        Upper bound: first term alone (overestimates)
        Lower bound: first two terms (subtracts the first correction)

    Parameters
    ----------
    N : int
        Number of terms already summed (tail starts at n = N+1).

    Returns
    -------
    tuple of (float, float)
        (lower_bound, upper_bound) for the tail Σ_{n>N} 1/L(2n).
    """
    phi_inv2 = 1.0 / (PHI ** 2)  # φ^{-2}
    phi_inv6 = phi_inv2 ** 3     # φ^{-6}

    # Leading geometric series: Σ_{n>N} φ^{-2n} = φ^{-2(N+1)} / (1 - φ^{-2})
    upper = phi_inv2 ** (N + 1) / (1.0 - phi_inv2)

    # First correction: -Σ_{n>N} φ^{-6n} = -φ^{-6(N+1)} / (1 - φ^{-6})
    correction = phi_inv6 ** (N + 1) / (1.0 - phi_inv6)

    lower = upper - correction

    return (lower, upper)


def theta_lucas_identity_precision(terms: int) -> tuple:
    """Compute the theta-Lucas identity with a rigorous error bound.

    The identity θ₃(1/φ²)² = 1 + 4·Σ_{n=1}^∞ 1/L(2n) is evaluated with
    a finite number of terms on both sides, and the truncation error is
    bounded analytically.

    The θ₃ series converges as q^{n²} (super-exponential), so its truncation
    error is negligible compared to the Lucas side. The dominant error comes
    from the Lucas reciprocal sum tail.

    Parameters
    ----------
    terms : int
        Number of terms for both θ₃ and the Lucas sum.

    Returns
    -------
    tuple of (float, float)
        (identity_value, error_bound) where:
        - identity_value = θ₃(q)² - (1 + 4·Σ_{n=1}^{terms} 1/L(2n))
          (should approach 0 as terms → ∞)
        - error_bound = rigorous upper bound on |Σ_{n>terms} 1/L(2n)|,
          scaled by 4 to match the identity
    """
    q = golden_nome()

    # θ₃ truncation error: 2·Σ_{n>terms} q^{n²} ≤ 2·q^{(terms+1)²}/(1-q)
    # This is negligible for terms ≥ 9 (< 1e-16)
    theta3_val = theta3_golden(terms)
    theta3_sq = theta3_val ** 2

    # θ₃² truncation error (via product rule and geometric bound)
    theta3_tail = 2.0 * q ** ((terms + 1) ** 2) / (1.0 - q)
    theta3_sq_error = 2.0 * abs(theta3_val) * theta3_tail + theta3_tail ** 2

    # Lucas side
    lrs = lucas_reciprocal_sum(terms)
    lucas_side = 1.0 + 4.0 * lrs

    # Lucas tail bound
    _, lucas_tail_upper = lucas_reciprocal_tail_bound(terms)

    # Total error bound: 4 * lucas_tail + theta3_sq_error
    total_error = 4.0 * lucas_tail_upper + theta3_sq_error

    identity_value = theta3_sq - lucas_side

    return (identity_value, total_error)


def convergence_rate_comparison(N: int) -> dict:
    """Compare convergence rates of the two sides of the theta-Lucas identity.

    θ₃(q)² = 1 + 4·Σ 1/L(2n), where q = 1/φ².

    Left side (θ₃ series): converges as q^{n²} — super-exponential.
        The n-th term is 2q^{n²}. Rate: -n²·ln(q) per term.

    Right side (Lucas sum): converges as φ^{-2n} — geometric.
        The n-th term is 1/L(2n) ≈ φ^{-2n}. Rate: -2n·ln(φ) per term.

    The θ₃ series wins dramatically: by n=9, q^{81} ≈ 10^{-34},
    while 1/L(18) ≈ φ^{-18} ≈ 10^{-4}.

    Parameters
    ----------
    N : int
        Evaluate rates at this term index.

    Returns
    -------
    dict
        Keys:
        - 'theta3_term_N': value of the N-th theta3 term (2·q^{N²})
        - 'lucas_term_N': value of 1/L(2N)
        - 'theta3_log_rate': -N²·log₁₀(q) (effective digits gained at term N)
        - 'lucas_log_rate': 2N·log₁₀(φ) (effective digits gained at term N)
        - 'theta3_terms_for_eps': terms needed for θ₃ truncation < machine epsilon
        - 'lucas_terms_for_eps': terms needed for Lucas truncation < machine epsilon
        - 'bottleneck': which side limits overall precision ('theta3' or 'lucas')
        - 'terms_for_1e15': terms needed for 1e-15 accuracy on the Lucas side
    """
    q = golden_nome()
    eps = 2.2e-16  # float64 machine epsilon

    # Term values at index N
    theta3_term = 2.0 * q ** (N * N)
    lucas_term = PHI ** (-2 * N)  # asymptotic approximation

    # Log rates (digits of accuracy per term)
    log10_q = math.log10(q)
    log10_phi = math.log10(PHI)
    theta3_log_rate = -N * N * log10_q  # total digits from N terms
    lucas_log_rate = 2 * N * log10_phi

    # Terms needed for machine epsilon
    # θ₃: q^{n²} < eps => n > sqrt(-log(eps)/log(1/q))
    log_inv_q = math.log(1.0 / q)
    theta3_terms_eps = math.ceil(math.sqrt(-math.log(eps) / log_inv_q))

    # Lucas: φ^{-2n} < eps => n > -log(eps)/(2·log(φ))
    lucas_terms_eps = math.ceil(-math.log(eps) / (2.0 * math.log(PHI)))

    # Terms for 1e-15
    lucas_terms_1e15 = math.ceil(-math.log(1e-15) / (2.0 * math.log(PHI)))

    bottleneck = 'lucas' if lucas_terms_eps > theta3_terms_eps else 'theta3'

    return {
        'theta3_term_N': theta3_term,
        'lucas_term_N': lucas_term,
        'theta3_log_rate': theta3_log_rate,
        'lucas_log_rate': lucas_log_rate,
        'theta3_terms_for_eps': theta3_terms_eps,
        'lucas_terms_for_eps': lucas_terms_eps,
        'bottleneck': bottleneck,
        'terms_for_1e15': lucas_terms_1e15,
    }
