"""Completed Riemann xi evaluation, both tiers.

Rigorous tier (flint ``acb``/``arb``): :func:`xi_from_s` and
:func:`centered_xi` are re-exported from :mod:`rh_verify.reciprocal_xi`,
which owns the canonical normalization

    xi(s) = 1/2 * s * (s - 1) * pi^(-s/2) * Gamma(s/2) * zeta(s).

Exploration tier (mpmath, ``_mp`` suffix): fast arbitrary-precision floats
with no error tracking.  Never cite an ``_mp`` value as a certificate.
"""

from __future__ import annotations

from collections.abc import Iterator
from contextlib import contextmanager

import mpmath as mp
from flint import acb, acb_series, arb, ctx

from .reciprocal_xi import centered_xi, xi_from_s

__all__ = [
    "centered_xi",
    "completed_lambda_mp",
    "log_xi_jet",
    "xi_from_s",
    "xi_jet",
    "xi_log_derivative_mp",
    "xi_mp",
    "xi_seed_mp",
]

_POLE_GUARD = 1e-8


def _series_length(degree: int) -> int:
    if not isinstance(degree, int) or degree < 0:
        raise ValueError("degree must be a nonnegative integer")
    return degree + 1


@contextmanager
def _series_cap(length: int) -> Iterator[None]:
    previous_cap = ctx.cap
    try:
        ctx.cap = max(previous_cap, length)
        yield
    finally:
        ctx.cap = previous_cap


def xi_jet(center: acb, degree: int) -> acb_series:
    """Rigorous Taylor jet of xi at ``center`` through ``degree``.

    Coefficient ``k`` of the returned series is ``xi^(k)(center) / k!``.
    """
    series_length = _series_length(degree)
    with _series_cap(series_length):
        variable = acb_series([center, acb(1)], series_length)
        return xi_from_s(variable)


def log_xi_jet(center: acb, degree: int) -> acb_series:
    """Rigorous Taylor jet of log xi at ``center`` through ``degree``.

    ``center`` must avoid the zeros of xi; flint reports failure by
    returning non-finite coefficients.
    """
    series_length = _series_length(degree)
    with _series_cap(series_length):
        return xi_jet(center, degree).log()


def xi_mp(s) -> mp.mpc:
    """Exploration-tier xi(s) via mpmath, with removable values at 0 and 1."""
    s = mp.mpmathify(s)
    if s == 0 or s == 1:
        return mp.mpf("0.5")
    return (
        mp.mpf("0.5")
        * s
        * (s - 1)
        * mp.power(mp.pi, -s / 2)
        * mp.gamma(s / 2)
        * mp.zeta(s)
    )


def xi_seed_mp(s) -> mp.mpc:
    """Exploration-tier xi prefactor 1/2 * s(s-1) * pi^(-s/2) * Gamma(s/2)."""
    s = mp.mpmathify(s)
    return (
        mp.mpf("0.5") * s * (s - 1) * mp.power(mp.pi, -s / 2) * mp.gamma(s / 2)
    )


def completed_lambda_mp(s) -> mp.mpc:
    """Exploration-tier completed zeta pi^(-s/2) * Gamma(s/2) * zeta(s)."""
    s = mp.mpmathify(s)
    return mp.power(mp.pi, -s / 2) * mp.gamma(s / 2) * mp.zeta(s)


def xi_log_derivative_mp(s) -> mp.mpc:
    """Exploration-tier (log xi)'(s) via the closed form.

    Uses 1/s + 1/(s-1) - log(pi)/2 + digamma(s/2)/2 + zeta'(s)/zeta(s);
    falls back to numerical differentiation near the removable points
    s = 0 and s = 1 where the closed form splits into singular pieces.
    """
    s = mp.mpmathify(s)
    if abs(s) < _POLE_GUARD or abs(s - 1) < _POLE_GUARD:
        return mp.diff(lambda w: mp.log(xi_mp(w)), s)
    return (
        1 / s
        + 1 / (s - 1)
        - mp.log(mp.pi) / 2
        + mp.digamma(s / 2) / 2
        + mp.zeta(s, derivative=1) / mp.zeta(s)
    )
