"""Rigorous reciprocal-xi evaluations shared by adjacent certificates.

The moment routine certifies

    2 * integral_0^infinity x^order / xi(1/2 + x) dx

for even nonnegative ``order``.  It uses a Taylor regularization across the
removable zeta pole at ``x = 1/2`` and a geometric tail bound obtained from
the completed-zeta recurrence over steps of length two.
"""

from __future__ import annotations

from dataclasses import dataclass
from fractions import Fraction
from functools import lru_cache
from math import factorial

from flint import acb, acb_series, arb

from .certificates import CertificationError, require_positive


REAL_CUTOFF = arb(80)
REGULARIZER_DEGREE = 48
RELATIVE_TOLERANCE = arb("1e-55")
REAL_ENDPOINTS = (
    "0.55",
    "1",
    "2",
    "4",
    "8",
    "12",
    "20",
    "32",
    "50",
    "65",
    "80",
)


@dataclass(frozen=True)
class ReciprocalXiMomentCertificate:
    """Interval data proving one even reciprocal-xi moment."""

    order: int
    compact_integral: arb
    pole_window_error: arb
    recurrence_growth_lower_bound: arb
    tail_step_ratio: arb
    tail_bound: arb
    enclosure: arb


def _as_arb(value: Fraction) -> arb:
    return arb(value.numerator) / value.denominator


@lru_cache(maxsize=None)
def _regularizer_coefficients(degree: int) -> tuple[Fraction, ...]:
    """Taylor coefficients of y/(1-exp(-y)) through ``degree``."""
    denominator = tuple(
        Fraction((-1) ** index, factorial(index + 1))
        for index in range(degree + 1)
    )
    coefficients = [Fraction(0) for _ in range(degree + 1)]
    coefficients[0] = Fraction(1)
    for order in range(1, degree + 1):
        coefficients[order] = -sum(
            denominator[index] * coefficients[order - index]
            for index in range(1, order + 1)
        )
    return tuple(coefficients)


def _regularizer_polynomial(value: acb) -> acb:
    result = acb(0)
    for coefficient in reversed(
        _regularizer_coefficients(REGULARIZER_DEGREE)
    ):
        result = result * value + acb(_as_arb(coefficient))
    return result


def xi_from_s(s: acb | acb_series) -> acb | acb_series:
    """Evaluate the completed Riemann xi function."""
    return (
        s
        * (s - 1)
        * (-s * acb.pi().log() / 2).exp()
        * (s / 2).gamma()
        * s.zeta()
        / 2
    )


def centered_xi(centered: acb | acb_series) -> acb | acb_series:
    """Evaluate xi(1/2 + centered)."""
    return xi_from_s(acb(arb(1) / 2) + centered)


def _regularized_centered_xi(centered: acb) -> acb:
    s = acb(arb(1) / 2) + centered
    logarithm_two = acb(2).log()
    y = (s - 1) * logarithm_two
    regularized_product = (
        s.dirichlet_eta()
        / logarithm_two
        * _regularizer_polynomial(y)
    )
    return (
        s
        * (-s * acb.pi().log() / 2).exp()
        * (s / 2).gamma()
        * regularized_product
        / 2
    )


def centered_log_xi_second_derivative() -> arb:
    """Certify (log xi(1/2+x))'' at x=0."""
    variable = acb_series([acb(arb(1) / 2), acb(1)], 3)
    logarithm = xi_from_s(variable).log()
    return 2 * arb(logarithm[2].real, abs(logarithm[2].imag))


def _regularizer_remainder() -> arb:
    cauchy_radius = arb(1) / 3
    window_radius = arb(2).log() / 20
    maximum = 1 / (2 - cauchy_radius.exp())
    ratio = window_radius / cauchy_radius
    return (
        maximum
        * ratio ** (REGULARIZER_DEGREE + 1)
        / (1 - ratio)
    )


def _pole_window_reciprocal_error() -> arb:
    remainder = _regularizer_remainder()
    s_box = arb(1, arb("0.05"))
    prefactor = (
        s_box
        / 2
        * (-s_box * arb.pi().log() / 2).exp()
        * (s_box / 2).gamma()
        / arb(2).log()
    )
    xi_error = abs(prefactor).upper() * remainder.upper()
    # The existing reciprocal-xi compact certificate proves xi(s)>9/20
    # throughout the real pole window 0.95 <= s <= 1.05.
    xi_lower = arb(9) / 20
    denominator = xi_lower * (xi_lower - xi_error)
    require_positive(denominator, "pole-window reciprocal denominator")
    return xi_error / denominator


def _zeta_tail_delta(real_part: arb) -> arb:
    return 2 ** (-real_part) * (1 + 2 / (real_part - 1))


def _recurrence_growth_lower_bound() -> arb:
    sigma = REAL_CUTOFF + arb(1) / 2
    rational = (
        (sigma + 2)
        * (sigma + 1)
        / (2 * arb.pi() * (sigma - 1))
    )
    zeta_ratio = (
        (1 - _zeta_tail_delta(sigma + 2))
        / (1 + _zeta_tail_delta(sigma))
    )
    return require_positive(
        rational * zeta_ratio,
        "reciprocal-xi recurrence growth",
    )


def _real_enclosure(value: acb) -> arb:
    return arb(value.real, abs(value.imag))


def _integrate_weight(
    order: int,
    left: str,
    right: str,
    *,
    regularized: bool = False,
) -> acb:
    xi_function = (
        _regularized_centered_xi if regularized else centered_xi
    )
    return acb.integral(
        lambda x, analytic: 2 * x**order / xi_function(x),
        arb(left),
        arb(right),
        rel_tol=RELATIVE_TOLERANCE,
    )


def certify_reciprocal_xi_even_moment(
    order: int,
) -> ReciprocalXiMomentCertificate:
    """Certify the full-line reciprocal-xi moment of even order."""
    if not isinstance(order, int) or order < 0 or order % 2:
        raise ValueError("order must be a nonnegative even integer")

    pieces = [
        _integrate_weight(order, "0", "0.45"),
        _integrate_weight(order, "0.45", "0.55", regularized=True),
    ]
    pieces.extend(
        _integrate_weight(order, left, right)
        for left, right in zip(REAL_ENDPOINTS, REAL_ENDPOINTS[1:])
    )
    compact = _real_enclosure(sum(pieces, acb(0)))

    pole_error = (
        arb(1) / 5
        * (arb(11) / 20) ** order
        * _pole_window_reciprocal_error()
    )

    first_tail = require_positive(
        _real_enclosure(_integrate_weight(order, "80", "82")),
        f"order-{order} first tail interval",
    )
    growth = _recurrence_growth_lower_bound()
    weight_ratio = ((REAL_CUTOFF + 2) / REAL_CUTOFF) ** order
    step_ratio = weight_ratio / growth
    if not step_ratio < 1:
        raise CertificationError(
            f"order-{order} tail ratio is not below one: {step_ratio}"
        )
    tail = require_positive(
        first_tail / (1 - step_ratio),
        f"order-{order} geometric tail",
    )

    enclosure = compact + arb(
        0,
        pole_error.upper() + tail.upper(),
    )
    require_positive(enclosure, f"order-{order} reciprocal-xi moment")
    return ReciprocalXiMomentCertificate(
        order=order,
        compact_integral=compact,
        pole_window_error=pole_error,
        recurrence_growth_lower_bound=growth,
        tail_step_ratio=step_ratio,
        tail_bound=tail,
        enclosure=enclosure,
    )
