#!/usr/bin/env python3
"""
L-Function Engine for BSDProver

Computes L(E,s) for elliptic curves using the Approximate Functional Equation (AFE),
which provides rapid convergence at s=1 unlike the raw Euler product.

The AFE for L(E,1) is:
    L(E,1) = 2 * sum_{n=1}^{X} (a_n / n) * exp(-2*pi*n / sqrt(N))

where N is the conductor, a_n are the Dirichlet coefficients computed via
multiplicativity from the Frobenius traces a_p, and X is chosen so that
the exponential decay makes the tail negligible.

Reference: Cremona, "Algorithms for Modular Elliptic Curves", Section 2.13
"""

import numpy as np
import mpmath
from decimal import Decimal, getcontext
import sympy as sp
from sympy.ntheory import factorint, legendre_symbol
from fractions import Fraction
from typing import Dict, List, Tuple, Optional, Union
from dataclasses import dataclass
import warnings
import math

# Set high precision
getcontext().prec = 200
mpmath.dps = 100

@dataclass
class LFunctionResult:
    """Result container for L-function computations"""
    value: complex
    derivative: Optional[complex] = None
    order_of_vanishing: int = 0
    precision: int = 50
    error_bound: Optional[float] = None
    computation_method: str = "approximate_functional_equation"
    primes_used: int = 0
    convergence_achieved: bool = True
    metadata: Dict = None

    def __post_init__(self):
        if self.metadata is None:
            self.metadata = {}

@dataclass
class PadicLFunctionResult:
    """Result container for p-adic L-function computations"""
    p: int
    value: complex
    precision: int
    error_bound: Optional[float] = None
    interpolation_points: List[complex] = None
    iwasawa_invariants: Dict = None
    metadata: Dict = None

    def __post_init__(self):
        if self.interpolation_points is None:
            self.interpolation_points = []
        if self.iwasawa_invariants is None:
            self.iwasawa_invariants = {}
        if self.metadata is None:
            self.metadata = {}


class LFunctionEngine:
    """
    L-function computation engine using the Approximate Functional Equation.

    The Euler product for L(E,s) converges only for Re(s) > 3/2, so it cannot
    be used directly at s=1. Instead, we use the AFE which converges rapidly
    due to the exponential damping factor exp(-2*pi*n/sqrt(N)).

    Reference: Cremona, "Algorithms for Modular Elliptic Curves", Ch. 2
    """

    def __init__(self, precision: int = 100, prime_limit: int = 100000):
        self.precision = precision
        self.prime_limit = prime_limit
        mpmath.dps = precision

        # Cache for computed values
        self._l_value_cache = {}
        self._ap_cache = {}

    def compute_l_function(self, curve, s: complex = 1, use_arb: bool = True) -> LFunctionResult:
        """
        Compute L(E,s) using the Approximate Functional Equation.

        Args:
            curve: EllipticCurve object with attributes a, b, conductor, discriminant
            s: Complex point to evaluate (default s=1 for BSD)
            use_arb: Ignored (kept for API compatibility)

        Returns:
            LFunctionResult with value and metadata
        """
        if curve.is_singular:
            raise ValueError("Cannot compute L-function for singular curve")

        # Check cache
        cache_key = (curve.curve_id, s, self.prime_limit)
        if cache_key in self._l_value_cache:
            return self._l_value_cache[cache_key]

        try:
            result = self._compute_l_function_afe(curve, s)
            self._l_value_cache[cache_key] = result
            return result
        except Exception as e:
            raise RuntimeError(f"L-function computation failed: {e}")

    def _compute_an_coefficients(self, curve, X: int) -> list:
        """
        Compute Dirichlet coefficients a_n for n=1..X using multiplicativity.

        The coefficients satisfy:
          a_1 = 1
          a_p = Frobenius trace (good primes) or 0/+-1 (bad primes)
          a_{p^k} = a_p * a_{p^{k-1}} - p * a_{p^{k-2}} (good primes)
          a_{p^k} = (a_p)^k (bad primes)
          a_{mn} = a_m * a_n when gcd(m,n) = 1

        We use a sieve-like approach: for each prime p, fill in a_{p^k},
        then compose via multiplicativity.
        """
        conductor = curve.conductor
        a_n = [0] * (X + 1)
        a_n[1] = 1

        # For each prime p <= X, compute a_{p^k} for all p^k <= X
        for p in sp.primerange(2, X + 1):
            is_bad = (conductor % p == 0)

            if is_bad:
                a_p = self._compute_bad_prime_ap(curve, p)
            else:
                a_p = self._compute_frobenius_trace(curve, p)

            # Fill a_{p^k} using the recurrence
            pk = p  # p^1
            a_pk_prev2 = 1   # a_{p^0} = 1
            a_pk_prev1 = a_p  # a_{p^1}
            a_n[pk] = a_p

            pk *= p  # p^2
            k = 2
            while pk <= X:
                if is_bad:
                    # Bad prime: a_{p^k} = (a_p)^k
                    a_pk = a_p ** k
                else:
                    # Good prime: a_{p^k} = a_p * a_{p^{k-1}} - p * a_{p^{k-2}}
                    a_pk = a_p * a_pk_prev1 - p * a_pk_prev2

                a_n[pk] = a_pk
                a_pk_prev2 = a_pk_prev1
                a_pk_prev1 = a_pk
                pk *= p
                k += 1

        # Now propagate via multiplicativity: a_{mn} = a_m * a_n for gcd(m,n)=1
        # We do this by iterating over primes and "spreading" the prime-power
        # contributions to composite indices.
        # Efficient approach: iterate n from 2 to X, decompose, multiply.
        # Since we already have a_{p^k} for all prime powers, we can build
        # a_n for composite n by factoring.

        for n in range(2, X + 1):
            if a_n[n] != 0:
                continue  # Already computed (prime power)

            # Factor n and compute a_n from prime power components
            factors = factorint(n)
            result = 1
            for p, e in factors.items():
                pe = p ** e
                if pe <= X and a_n[pe] != 0:
                    result *= a_n[pe]
                elif e == 1 and p <= X:
                    result *= a_n[p]
                else:
                    # Need to compute a_{p^e} from recurrence
                    a_pe = self._compute_a_prime_power(curve, p, e)
                    result *= a_pe

            a_n[n] = result

        return a_n

    def _compute_a_prime_power(self, curve, p: int, e: int) -> int:
        """Compute a_{p^e} using the recurrence relation."""
        conductor = curve.conductor
        is_bad = (conductor % p == 0)

        if is_bad:
            a_p = self._compute_bad_prime_ap(curve, p)
            return a_p ** e

        a_p = self._compute_frobenius_trace(curve, p)
        if e == 0:
            return 1
        if e == 1:
            return a_p

        # Good prime recurrence: a_{p^k} = a_p * a_{p^{k-1}} - p * a_{p^{k-2}}
        prev2 = 1    # a_{p^0}
        prev1 = a_p  # a_{p^1}
        for _ in range(2, e + 1):
            curr = a_p * prev1 - p * prev2
            prev2 = prev1
            prev1 = curr
        return prev1

    def _compute_l_function_afe(self, curve, s: complex) -> LFunctionResult:
        """
        Compute L(E,s) using the Approximate Functional Equation.

        At s=1:
            L(E,1) = 2 * sum_{n=1}^{X} (a_n/n) * exp(-2*pi*n/sqrt(N))

        The exponential decay ensures rapid convergence. For conductor N,
        terms with n >> sqrt(N)/(2*pi) are negligible.

        Reference: Cremona, "Algorithms for Modular Elliptic Curves", 2.13
        """
        mpmath.mp.dps = self.precision
        conductor = curve.conductor

        if conductor <= 0:
            raise ValueError(f"Invalid conductor: {conductor}")

        sqrt_N = mpmath.sqrt(conductor)

        # Cutoff: ensure enough terms for convergence
        # exp(-2*pi*X/sqrt(N)) should be < 10^{-precision}
        # => X > sqrt(N)/(2*pi) * precision * ln(10)
        min_terms = int(float(sqrt_N / (2 * mpmath.pi)) * self.precision * math.log(10)) + 50
        X = max(200, min(min_terms, 50000))  # Cap at 50000 to avoid memory issues

        # Compute Dirichlet coefficients
        a_n = self._compute_an_coefficients(curve, X)

        # Evaluate AFE sum
        L_value = mpmath.mpf(0)
        two_pi_over_sqrt_N = 2 * mpmath.pi / sqrt_N

        for n in range(1, X + 1):
            if a_n[n] == 0:
                continue
            term = mpmath.mpf(a_n[n]) / n * mpmath.exp(-two_pi_over_sqrt_N * n)
            L_value += term

        L_value *= 2

        # Error bound from truncation: remaining terms bounded by
        # 2 * sum_{n>X} |a_n|/n * exp(-2*pi*n/sqrt(N))
        # Since |a_n| <= d(n)*sqrt(n) (Ramanujan-Petersson conjecture, proved by Deligne),
        # and the exponential decay dominates, the tail is bounded by:
        tail_bound = float(2 * X * mpmath.exp(-two_pi_over_sqrt_N * X))

        # Detect order of vanishing
        order_vanishing = 0
        if abs(s - 1) < 1e-10:
            if abs(float(L_value)) < 10 ** (-(self.precision // 4)):
                order_vanishing = 1

        primes_used = len(list(sp.primerange(2, X + 1)))

        return LFunctionResult(
            value=complex(L_value),
            order_of_vanishing=order_vanishing,
            precision=self.precision,
            error_bound=tail_bound,
            computation_method="approximate_functional_equation",
            primes_used=primes_used,
            convergence_achieved=True,
            metadata={
                "conductor": conductor,
                "evaluation_point": s,
                "cutoff_X": X,
                "sqrt_N": float(sqrt_N),
            }
        )

    def _compute_frobenius_trace(self, curve, p: int) -> int:
        """
        Compute trace of Frobenius a_p = p + 1 - #E(F_p) for good prime p.

        This is the standard point-counting algorithm over F_p.
        """
        cache_key = (curve.curve_id, p)
        if cache_key in self._ap_cache:
            return self._ap_cache[cache_key]

        a, b = int(curve.a), int(curve.b)
        count = 1  # Point at infinity

        if p == 2:
            # Over F_2, every element is a perfect square (0^2=0, 1^2=1),
            # so each x gives exactly one y-solution. a_2 = 0 for all curves.
            for x in range(2):
                count += 1
        else:
            for x in range(p):
                y_squared = (pow(x, 3, p) + a * x + b) % p

                if y_squared == 0:
                    count += 1
                else:
                    if legendre_symbol(y_squared, p) == 1:
                        count += 2

        a_p = p + 1 - count
        self._ap_cache[cache_key] = a_p
        return a_p

    def _compute_bad_prime_ap(self, curve, p: int) -> int:
        """
        Compute a_p for bad prime p | N (conductor).

        - Additive reduction (p^2 | discriminant): a_p = 0
        - Split multiplicative reduction: a_p = 1
        - Non-split multiplicative reduction: a_p = -1

        The split/non-split distinction uses the Legendre symbol of the
        discriminant modulo p for p >= 3.
        """
        discriminant = int(curve.discriminant)

        if discriminant % (p ** 2) == 0:
            return 0  # Additive reduction

        # Multiplicative reduction
        if p == 2:
            return 1 if discriminant % 8 in [1, 7] else -1
        else:
            return legendre_symbol(discriminant, p)

    def _compute_local_factor(self, a_p: int, p: int, s: complex, is_bad: bool) -> complex:
        """Compute local Euler factor (kept for functional equation verification)."""
        if is_bad:
            if a_p == 0:
                return 1
            else:
                return 1 - a_p * (p ** (-s))
        else:
            p_to_neg_s = p ** (-s)
            p_to_neg_2s = p ** (-2 * s)
            return 1 - a_p * p_to_neg_s + p_to_neg_2s

    def compute_l_derivative(self, curve, s: complex = 1) -> complex:
        """
        Compute L'(E,1) using the AFE for root number w = -1.

        When w = -1, L(E,1) = 0 and differentiating the AFE at s=1 gives:

          L'(E,1) = 2 * sum_{n=1}^X (a_n / n) * E_1(2*pi*n / sqrt(N))

        where E_1(x) = integral_x^inf e^{-t}/t dt is the exponential integral.

        All correction terms from the gamma/log factors cancel exactly due
        to the symmetry of the functional equation at s=1.

        Reference: Cohen, "A Course in Computational Algebraic Number Theory",
                   Section 7.5.3; Cremona, Section 2.13
        """
        mpmath.mp.dps = self.precision
        conductor = curve.conductor

        if conductor <= 0:
            raise ValueError(f"Invalid conductor: {conductor}")

        sqrt_N = mpmath.sqrt(conductor)
        c = 2 * mpmath.pi / sqrt_N

        # Cutoff: same as for L(E,1)
        min_terms = int(float(sqrt_N / (2 * mpmath.pi)) * self.precision * math.log(10)) + 50
        X = max(200, min(min_terms, 50000))

        a_n = self._compute_an_coefficients(curve, X)

        L_prime = mpmath.mpf(0)

        for n in range(1, X + 1):
            if a_n[n] == 0:
                continue

            cn = c * n
            e1_cn = mpmath.e1(cn)

            term = mpmath.mpf(a_n[n]) / n * e1_cn
            L_prime += term

        L_prime *= 2

        return complex(L_prime)

    def verify_functional_equation(self, curve, s1: complex = 0.5 + 3j) -> bool:
        """
        Verify the functional equation Lambda(s) = epsilon * Lambda(2-s).

        Uses the Euler product (which converges for Re(s) > 3/2) to check
        consistency at test points.
        """
        try:
            conductor = curve.conductor

            # Compute L(E, s1) and L(E, 2-s1) via Euler product (converges for large Re(s))
            # For verification, use points where the Euler product converges
            s_test = complex(2.0, 1.0)  # Re(s) = 2 > 3/2
            s_test_mirror = 2 - s_test

            L1 = self._compute_euler_product_at_s(curve, s_test)
            L2 = self._compute_euler_product_at_s(curve, s_test_mirror)

            Lambda1 = (conductor ** (s_test / 2) * (2 * np.pi) ** (-s_test) *
                       complex(mpmath.gamma(s_test)) * L1)
            Lambda2 = (conductor ** (s_test_mirror / 2) * (2 * np.pi) ** (-s_test_mirror) *
                       complex(mpmath.gamma(s_test_mirror)) * L2)

            ratio = abs(Lambda1 / Lambda2) if abs(Lambda2) > 1e-50 else float('inf')
            return abs(ratio - 1) < 0.01 or abs(ratio + 1) < 0.01

        except Exception:
            return False

    def _compute_euler_product_at_s(self, curve, s: complex) -> complex:
        """Euler product computation for Re(s) > 3/2 (verification purposes only)."""
        conductor = curve.conductor
        L_product = mpmath.mpc(1, 0)

        for p in sp.primerange(2, min(self.prime_limit, 10000) + 1):
            if conductor % p == 0:
                a_p = self._compute_bad_prime_ap(curve, p)
            else:
                a_p = self._compute_frobenius_trace(curve, p)

            local_factor = self._compute_local_factor(a_p, p, s, conductor % p == 0)
            if abs(local_factor) > 1e-50:
                L_product /= local_factor

        return complex(L_product)


class PadicLFunction:
    """
    p-adic L-function framework using Mazur-Tate-Teitelbaum construction.

    Note: This provides a simplified implementation. Full p-adic L-function
    computation requires specialized p-adic arithmetic libraries.
    """

    def __init__(self, p: int, precision: int = 50):
        if not sp.isprime(p):
            raise ValueError(f"p={p} must be prime")

        self.p = p
        self.precision = precision
        self.classical_engine = LFunctionEngine(precision=precision)

    def compute_p_adic_l_function(self, curve, s: complex = 1) -> PadicLFunctionResult:
        """
        Compute p-adic L-function L_p(E,s) via interpolation of classical values.

        Note: This is a simplified implementation using Lagrange interpolation
        of classical L-values. A rigorous implementation would require full
        p-adic arithmetic and overconvergent modular symbols.
        """
        if curve.conductor % self.p == 0:
            raise ValueError(f"Cannot compute L_{self.p} for curve with bad reduction at {self.p}")

        try:
            interpolation_points = self._generate_interpolation_points(s)
            classical_values = []

            for point in interpolation_points:
                classical_result = self.classical_engine.compute_l_function(curve, point)
                classical_values.append(classical_result.value)

            p_adic_value = self._interpolate_p_adic(interpolation_points, classical_values, s)

            return PadicLFunctionResult(
                p=self.p,
                value=p_adic_value,
                precision=self.precision,
                interpolation_points=interpolation_points,
                iwasawa_invariants={
                    "lambda": 0, "mu": 0, "nu": 0,
                    "note": "placeholder -- requires specialized p-adic computation"
                },
                metadata={
                    "conductor": curve.conductor,
                    "interpolation_method": "lagrange_classical"
                }
            )

        except Exception as e:
            raise RuntimeError(f"p-adic L-function computation failed: {e}")

    def _generate_interpolation_points(self, s: complex) -> List[complex]:
        """Generate interpolation points s = 1 + k*(p-1) for k in range."""
        points = []
        for k in range(-5, 6):
            point = 1 + k * (self.p - 1)
            points.append(complex(point))
        return points

    def _interpolate_p_adic(self, points: List[complex], values: List[complex],
                           target: complex) -> complex:
        """Lagrange interpolation (simplified; true p-adic interpolation requires p-adic arithmetic)."""
        result = 0
        n = len(points)
        for i in range(n):
            basis = 1
            for j in range(n):
                if i != j:
                    basis *= (target - points[j]) / (points[i] - points[j])
            result += values[i] * basis
        return result