"""Certified integration and tail-bound primitives.

Everything here returns rigorous ``arb`` enclosures.  There is deliberately
no mpmath quadrature wrapper: ``mp.quad`` output carries no error bound and
must never be dressed as a certificate.  Use it freely in scratch, then
reconstruct load-bearing integrals through :func:`certified_real_integral`.
"""

from __future__ import annotations

from math import factorial

from flint import acb, arb

from .certificates import CertificationError

__all__ = [
    "certified_piecewise_real_integral",
    "certified_real_integral",
    "geometric_tail",
    "real_enclosure",
    "taylor_remainder",
]

RELATIVE_TOLERANCE = arb("1e-55")


def real_enclosure(value: acb) -> arb:
    """Fold a complex ball into a real enclosure of its real part.

    The imaginary radius is added to the real radius, so the result
    contains ``Re(z)`` for every ``z`` in the input ball even when the
    computation leaked a spurious imaginary component.
    """
    return arb(value.real, abs(value.imag))


def certified_real_integral(
    integrand,
    left: arb,
    right: arb,
    *,
    rel_tol: arb = RELATIVE_TOLERANCE,
    label: str = "integral",
) -> arb:
    """Certified enclosure of a real integral of an analytic integrand.

    ``integrand`` is called as ``integrand(z)`` with ``z`` an ``acb`` and
    must be analytic on the integration path.  Singularities inside the
    path silently poison the enclosure; split the path at them instead.
    """
    value = acb.integral(
        lambda z, analytic: integrand(z),
        arb(left),
        arb(right),
        rel_tol=rel_tol,
    )
    enclosure = real_enclosure(value)
    if not enclosure.is_finite():
        raise CertificationError(f"{label} enclosure is not finite: {enclosure}")
    return enclosure


def certified_piecewise_real_integral(
    integrand,
    endpoints,
    *,
    rel_tol: arb = RELATIVE_TOLERANCE,
    label: str = "integral",
) -> arb:
    """Certified integral over consecutive pieces of an endpoint partition.

    ``endpoints`` is an increasing sequence; each adjacent pair becomes one
    :func:`certified_real_integral` piece.  Place a breakpoint at every
    point where analyticity or smoothness degrades.
    """
    points = [arb(point) for point in endpoints]
    if len(points) < 2:
        raise CertificationError(f"{label} needs at least two endpoints")
    total = arb(0)
    for index, (piece_left, piece_right) in enumerate(zip(points, points[1:])):
        if not piece_left < piece_right:
            raise CertificationError(
                f"{label} endpoints not increasing at piece {index}: "
                f"{piece_left} versus {piece_right}"
            )
        total += certified_real_integral(
            integrand,
            piece_left,
            piece_right,
            rel_tol=rel_tol,
            label=f"{label} piece {index}",
        )
    return total


def geometric_tail(
    first_omitted: arb, step_ratio: arb, *, label: str = "tail"
) -> arb:
    """Bound ``sum_k first_omitted * step_ratio^k`` for a proven ratio.

    ``step_ratio`` must be provably inside ``[0, 1)``; the caller owns the
    proof that it really dominates every consecutive-term ratio of the
    series being bounded.
    """
    if not step_ratio >= 0 or not step_ratio < 1:
        raise CertificationError(
            f"{label} step ratio is not provably in [0, 1): {step_ratio}"
        )
    return first_omitted / (1 - step_ratio)


def taylor_remainder(derivative_bound: arb, radius: arb, degree: int) -> arb:
    """Lagrange remainder bound ``M * r^(n+1) / (n+1)!`` for degree ``n``.

    ``derivative_bound`` must dominate ``|f^(n+1)|`` on the whole closed
    disk of the given radius around the expansion center.
    """
    if degree < 0:
        raise CertificationError("Taylor remainder needs a nonnegative degree")
    return (
        abs(derivative_bound)
        * abs(radius) ** (degree + 1)
        / factorial(degree + 1)
    )
