#!/usr/bin/env python3
"""
Certified L-function engine (R1)

Every output is an Arb ball whose radius includes a PROVEN truncation
bound, so `L_leading(r)` is a mathematical enclosure of L^(r)(E,1)/r!.

Formulas (derived from the integral representation of the completed
L-function Lambda(s) = (sqrt(N)/2pi)^s Gamma(s) L(s), functional equation
Lambda(s) = eps * Lambda(2-s), eps = +/-1):

    L^(r)(E,1)/r! = 2 * sum_{n>=1} (a_n/n) * G_r(2*pi*n/sqrt(N)),

valid when ord_{s=1} L >= r and eps = (-1)^r, with G_r from
certified_arithmetic (G_0 = e^{-x}, G_1 = E_1, ...).

Tail bound (PROVEN): |a_n| <= d(n) sqrt(n) < 2n  (Deligne's bound for
weight-2 newforms + the elementary d(n) < 2 sqrt(n)), and
G_r(x) <= e^{-x} for r=0,  G_r(x) <= e^{-x}/x for r=1,2,3. Hence with
q = e^{-c}, c = 2pi/sqrt(N):

    |tail_0| <= 4 * q^{X+1}/(1-q)
    |tail_r| <= (4/(cX)) * q^{X+1}/(1-q),  r >= 1.

Root number (certified, and an independent CONDUCTOR VALIDATOR):
for f(tau) = sum a_n q^n the modular form of level N attached to E
(modularity theorem), A(t) := f(it/sqrt(N)) satisfies

    A(1/t) = eps * t^2 * A(t)    for all t > 0.

We evaluate both sides as balls at t = 6/5; the ratio ball must contain
exactly one of {+1, -1}. If it contains neither, the conductor fed in is
WRONG -- this is the FCE class-3 'functional equation as boundary
reflection' check, and it independently validates Tate's algorithm at
p = 2, 3 for every curve.
"""

from dataclasses import dataclass
from fractions import Fraction
from typing import Dict, List, Optional

import numpy as np
import sympy as sp
from flint import arb

from . import certified_arithmetic as ca
from .tate_algorithm import GlobalReductionData, analyze_reduction


class ConductorInvalidError(RuntimeError):
    """Functional-equation ball contained neither +1 nor -1: the conductor
    (hence Tate's algorithm output) is inconsistent with modularity."""


class RankBudgetExceeded(RuntimeError):
    """Rank >= 2 detected but deep generators (G_2/G_3 rigorous integrals)
    were disabled by the caller's budget (e.g. large sweeps)."""


@dataclass
class RankCertificate:
    epsilon: int                 # certified functional-equation sign
    rank: int                    # analytic rank (see `certified`)
    leading_coefficient: arb     # ball enclosing L^(rank)(E,1)/rank!
    certified: bool              # True: rank fully certified (0 or 1)
    conditional_on: str          # assumption for rank 2/3
    L1_ball: Optional[arb] = None


class CertifiedLFunction:
    def __init__(self, a: int, b: int, dps: int = 30,
                 reduction: Optional[GlobalReductionData] = None):
        self.a, self.b = int(a), int(b)
        self.dps = dps
        ca.set_precision(dps)
        self.g = reduction if reduction is not None else analyze_reduction(a, b)
        self.N = self.g.conductor
        self.disc_short = -16 * (4 * self.a ** 3 + 27 * self.b ** 2)
        self._ap_cache: Dict[int, int] = {}
        self._an_cache: Dict[int, List[int]] = {}
        self._eps: Optional[int] = None

    # ------------------------------------------------------------------ #
    # Exact Dirichlet coefficients
    # ------------------------------------------------------------------ #

    def a_p(self, p: int) -> int:
        """Exact a_p: Tate data at bad primes, point count at good primes
        (on the minimal general model when the short model is singular
        mod p, i.e. p | u)."""
        hit = self._ap_cache.get(p)
        if hit is not None:
            return hit

        if self.N % p == 0:
            ap = self.g.local[p].a_p
        elif self.disc_short % p == 0:
            ap = self._ap_general_model(self.g.minimal.ainvs, p)
        else:
            ap = self._ap_short_numpy(self.a, self.b, p)
        self._ap_cache[p] = ap
        return ap

    @staticmethod
    def _ap_short_numpy(a: int, b: int, p: int) -> int:
        """a_p on y^2 = x^3 + ax + b via vectorized exact point count."""
        if p == 2:
            # nonsingular mod 2 short model: each x has exactly one y
            return 0
        x = np.arange(p, dtype=np.int64)
        fx = (x * x % p * x + (a % p) * x + b) % p
        qr = np.zeros(p, dtype=np.int8)
        qr[(x * x) % p] = 1
        chi = np.where(fx == 0, 0, np.where(qr[fx] == 1, 1, -1))
        return int(-chi.sum())

    @staticmethod
    def _ap_general_model(ainvs, p: int) -> int:
        """a_p on a general model (used when p | u): exact, O(p)."""
        a1, a2, a3, a4, a6 = ainvs
        if p == 2:
            count = 1
            for x in range(2):
                for y in range(2):
                    if (y * y + a1 * x * y + a3 * y) % 2 == \
                       (x ** 3 + a2 * x * x + a4 * x + a6) % 2:
                        count += 1
            return 2 + 1 - count
        total = 0
        for x in range(p):
            fx = (x ** 3 + a2 * x * x + a4 * x + a6) % p
            d = ((a1 * x + a3) ** 2 + 4 * fx) % p
            if d != 0:
                total += 1 if pow(d, (p - 1) // 2, p) == 1 else -1
        return -total

    def a_n_list(self, X: int) -> List[int]:
        """Exact a_n for n = 1..X via multiplicativity (spf sieve)."""
        hit = self._an_cache.get(X)
        if hit is not None:
            return hit

        a = [0] * (X + 1)
        a[1] = 1
        # smallest-prime-factor sieve
        spf = np.zeros(X + 1, dtype=np.int64)
        for p in sp.primerange(2, X + 1):
            sl = spf[p::p]
            sl[sl == 0] = p
            spf[p::p] = sl

        # prime powers
        for p in sp.primerange(2, X + 1):
            ap = self.a_p(p)
            is_bad = (self.N % p == 0)
            a[p] = ap
            pk, prev2, prev1, k = p * p, 1, ap, 2
            while pk <= X:
                if is_bad:
                    apk = ap ** k
                else:
                    apk = ap * prev1 - p * prev2
                a[pk] = apk
                prev2, prev1 = prev1, apk
                pk *= p
                k += 1

        # composites via multiplicativity
        for n in range(2, X + 1):
            if a[n] != 0 or n == 1:
                continue
            p = int(spf[n])
            pe, m = p, n // p
            while m % p == 0:
                pe *= p
                m //= p
            if m == 1:
                continue  # pure prime power, already set (may be 0 genuinely)
            a[n] = a[pe] * a[m]

        # note: a prime power with a[pe] = 0 legitimately gives a[n] = 0;
        # the loop above skips n with a[n] != 0 only as an optimization --
        # composites whose true value is 0 are correctly 0 by init.
        # But composites n = pe*m with a[pe] != 0, a[m] != 0 that were
        # skipped because a[n] != 0 cannot occur (a[n] starts at 0).
        self._an_cache[X] = a
        return a

    # ------------------------------------------------------------------ #
    # Cutoff and certified tails
    # ------------------------------------------------------------------ #

    def _cutoff(self, scale: float = 1.0) -> int:
        import math
        sqrtN = math.sqrt(self.N)
        X = int(sqrtN / (2 * math.pi) * (self.dps + 6) * math.log(10) * scale) + 30
        return min(max(X, 60), 400_000)

    def _tail_bound(self, r: int, X: int, c: arb) -> arb:
        """Proven bound on |2 sum_{n>X} (a_n/n) G_r(c n)| (see module doc)."""
        q = (-c).exp()
        geo = q ** (X + 1) / (1 - q)
        if r == 0:
            return 4 * geo
        return 4 / (c * X) * geo

    # ------------------------------------------------------------------ #
    # Leading coefficients
    # ------------------------------------------------------------------ #

    def L_leading(self, r: int) -> arb:
        """
        Certified ball for L^(r)(E,1)/r!  (valid as the leading Taylor
        coefficient when ord_{s=1} L >= r and eps = (-1)^r).
        """
        ca.set_precision(self.dps)
        X = self._cutoff()
        an = self.a_n_list(X)
        c = 2 * arb.pi() / arb(self.N).sqrt()

        total = arb(0)
        if r == 0:
            q = (-c).exp()
            qn = arb(1)
            for n in range(1, X + 1):
                qn *= q
                if an[n] == 0:
                    continue
                total += arb(an[n]) / n * qn
        else:
            for n in range(1, X + 1):
                if an[n] == 0:
                    continue
                x_n = c * n
                # negligible-term shortcut: bound the term instead of
                # evaluating G_r when e^{-x_n} is already below tolerance
                g = ca.G_r(r, x_n, self.dps,
                           cache_key=("G", self.N, r, n))
                total += arb(an[n]) / n * g
        total *= 2
        tail = self._tail_bound(r, X, c)
        return ca.add_error(total, tail.abs_upper())

    # ------------------------------------------------------------------ #
    # Certified root number (and conductor validation)
    # ------------------------------------------------------------------ #

    def _A_ball(self, t_num: int, t_den: int) -> arb:
        """A(t) = sum a_n exp(-2 pi n t / sqrt(N)) with certified tail."""
        import math
        t = Fraction(t_num, t_den)
        sqrtN = math.sqrt(self.N)
        X = int(sqrtN / (2 * math.pi) / float(t) * (self.dps + 8) * math.log(10)) + 30
        X = min(max(X, 60), 400_000)
        an = self.a_n_list(X)
        c = 2 * arb.pi() / arb(self.N).sqrt() * t_num / t_den
        q = (-c).exp()
        qn = arb(1)
        total = arb(0)
        for n in range(1, X + 1):
            qn *= q
            if an[n] == 0:
                continue
            total += arb(an[n]) * qn
        # tail: |a_n| < 2n; sum_{n>X} n q^n <= (X+1) q^{X+1}/(1-q)^2
        tail = 2 * (X + 1) * q ** (X + 1) / (1 - q) ** 2
        return ca.add_error(total, tail.abs_upper())

    def root_number(self) -> int:
        """
        Certified eps via the Fricke involution: A(1/t) = eps t^2 A(t).
        Doubles as a conductor validator (raises ConductorInvalidError).
        """
        if self._eps is not None:
            return self._eps
        ca.set_precision(self.dps)
        for (tn, td) in ((6, 5), (4, 3), (3, 2)):
            At = self._A_ball(tn, td)
            if ca.contains_zero(At):
                continue  # accidental near-zero; try another t
            Ainv = self._A_ball(td, tn)
            ratio = Ainv / (arb(tn * tn) / (td * td) * At)
            has_p = ca.contains_value(ratio, 1)
            has_m = ca.contains_value(ratio, -1)
            if has_p and not has_m:
                self._eps = 1
                return 1
            if has_m and not has_p:
                self._eps = -1
                return -1
            if not has_p and not has_m:
                raise ConductorInvalidError(
                    f"A(1/t)/(t^2 A(t)) = {ratio} contains neither +1 nor -1: "
                    f"conductor N={self.N} is inconsistent with modularity")
            # contains both: not enough precision; try next t
        raise ConductorInvalidError(
            f"root number undecidable at dps={self.dps} for N={self.N}; "
            "increase precision")

    # ------------------------------------------------------------------ #
    # Certified analytic rank
    # ------------------------------------------------------------------ #

    def analytic_rank(self, max_r: int = 3) -> RankCertificate:
        """
        Certified rank logic:
          eps=+1: L(1) ball excludes 0 -> rank 0 (certified).
                  contains 0 -> compute L''(1)/2!; nonzero -> rank 2
                  (conditional on L(1) = 0 exactly).
          eps=-1: L(1) = 0 is a THEOREM (odd functional equation).
                  L'(1) ball excludes 0 -> rank 1 (certified).
                  contains 0 -> L'''(1)/3! nonzero -> rank 3 (conditional).
        """
        eps = self.root_number()

        if eps == 1:
            L0 = self.L_leading(0)
            sign = ca.certified_sign(L0)
            if sign == 1:
                return RankCertificate(1, 0, L0, True, "", L0)
            if sign == -1:
                raise ArithmeticError(
                    f"L(E,1) certified negative ({L0}) -- impossible; "
                    "computation is inconsistent")
            if max_r < 2:
                raise RankBudgetExceeded(
                    f"eps=+1 with L(1) straddling 0 (rank >= 2); deep "
                    f"generators disabled (max_r={max_r})")
            L2 = self.L_leading(2)
            if ca.certified_sign(L2) is not None:
                return RankCertificate(
                    1, 2, L2, False,
                    f"assumes L(E,1)=0 exactly (|L(1)| < {ca.width(L0):.2e} "
                    "certified; exact vanishing not numerically certifiable)",
                    L0)
            raise ArithmeticError(
                "rank undecided: eps=+1, both L(1) and L''(1)/2 straddle 0 "
                f"at dps={self.dps} (rank >= 4?)")

        # eps = -1: L(1) = 0 exactly (theorem)
        L1 = self.L_leading(1)
        if ca.certified_sign(L1) is not None:
            return RankCertificate(-1, 1, L1, True, "", arb(0))
        if max_r < 3:
            raise RankBudgetExceeded(
                f"eps=-1 with L'(1) straddling 0 (rank >= 3); deep "
                f"generators disabled (max_r={max_r})")
        L3 = self.L_leading(3)
        if ca.certified_sign(L3) is not None:
            return RankCertificate(
                -1, 3, L3, False,
                f"assumes L'(E,1)=0 exactly (|L'(1)| < {ca.width(L1):.2e} "
                "certified)", arb(0))
        raise ArithmeticError(
            "rank undecided: eps=-1, both L'(1) and L'''(1)/6 straddle 0 "
            f"at dps={self.dps} (rank >= 5?)")