#!/usr/bin/env python3
"""Reproduce the numerical tables on the verified-numerics book page.

Requirements
------------
CPython 3.10+ and mpmath 1.3.0.

The script performs three deterministic audits:

1. Bessel J_0 root errors from backward order recurrence at several cutoffs
   and three terminal-tail orders.
2. Exact-rational alternating-series enclosures at the two endpoints of the
   10^-49 bracket for j_{0,1}.
3. Schwarzschild s=2, ell=2 continued-fraction roots with the large-order
   tail derived in Chapter 4, plus cutoff, tail-order, and precision sweeps.

All decimal constants are parsed from strings. The Bessel certificate uses
fractions.Fraction and therefore does not depend on floating-point rounding.
The recurrence tables are convergence studies, not rigorous enclosures.
"""

from __future__ import annotations

import argparse
import platform
from decimal import Decimal, localcontext
from fractions import Fraction
from math import factorial

import mpmath as mp


BESSEL_LEFT = "2.4048255576957727686216318793264546431242449091459"
BESSEL_RIGHT = "2.4048255576957727686216318793264546431242449091460"
BESSEL_CUTOFFS = (4, 8, 12, 16, 20, 24)
BESSEL_DPS = 100

QNM_CUTOFFS = (50, 100, 200, 400)
QNM_DPS = 80
QNM_TAIL_ORDER = 4
QNM_INITIAL_OMEGA_RE = "0.75"
QNM_INITIAL_OMEGA_IM = "-0.18"
QNM_REFERENCE_CUTOFF = 3200


def decimal_fraction(text: str) -> Fraction:
    """Parse a finite decimal string exactly."""
    sign = -1 if text.startswith("-") else 1
    unsigned = text.lstrip("+-")
    whole, fractional = unsigned.split(".")
    numerator = int(whole + fractional)
    return sign * Fraction(numerator, 10 ** len(fractional))


def fraction_decimal(value: Fraction, digits: int = 150) -> Decimal:
    """Convert a rational to Decimal for display only."""
    with localcontext() as context:
        context.prec = digits
        return Decimal(value.numerator) / Decimal(value.denominator)


def bessel_terminal_ratio(z: mp.mpf, cutoff: int, order: int) -> mp.mpf:
    """Approximate R_N=J_{N+1}(z)/J_N(z) at the terminal row."""
    if order == 0:
        return mp.mpf("0")

    leading = z / (2 * (cutoff + 1))
    if order == 1:
        return leading
    if order == 2:
        correction = z**2 / (4 * (cutoff + 1) * (cutoff + 2))
        return leading * (1 + correction)
    raise ValueError("Bessel tail order must be 0, 1, or 2.")


def bessel_residual(z: mp.mpf, cutoff: int, order: int) -> mp.mpf:
    """Evaluate H_N(z)=2/z-R_1 by backward recurrence."""
    ratio = bessel_terminal_ratio(z, cutoff, order)
    for n in range(cutoff, 1, -1):
        ratio = 1 / (2 * n / z - ratio)
    return 2 / z - ratio


def bessel_root(cutoff: int, order: int) -> mp.mpf:
    """Solve the finite Bessel recurrence residual near j_{0,1}."""
    tolerance = mp.power(10, -(mp.mp.dps - 15))
    return mp.findroot(
        lambda z: bessel_residual(z, cutoff, order),
        (mp.mpf("2.3"), mp.mpf("2.5")),
        tol=tolerance,
        maxsteps=50,
    )


def bessel_rational_enclosure(
    x: Fraction,
    last_index: int = 50,
) -> tuple[Fraction, Fraction, Fraction]:
    """Return S_50-a_51 < J_0(x) < S_50 and the next-term magnitude."""
    t = x * x / 4
    term = Fraction(1)
    partial = term
    for k in range(1, last_index + 1):
        term *= -t / (k * k)
        partial += term

    next_term = term * (-t) / ((last_index + 1) ** 2)
    if next_term >= 0:
        raise AssertionError("The selected truncation must have a negative next term.")
    return partial + next_term, partial, -next_term


def print_bessel_audit() -> None:
    """Print the Bessel convergence table and exact bracket."""
    mp.mp.dps = BESSEL_DPS
    reference = mp.besseljzero(0, 1)

    print("\nBessel backward-recurrence audit")
    print(f"working decimal digits: {BESSEL_DPS}")
    print("root solver: complex/real secant, starts 2.3 and 2.5")
    print("N         K=0 error         K=1 error         K=2 error")
    for cutoff in BESSEL_CUTOFFS:
        errors = [
            abs(bessel_root(cutoff, order) - reference)
            for order in range(3)
        ]
        print(
            f"{cutoff:2d}  "
            + "  ".join(f"{mp.nstr(error, 8):>16}" for error in errors)
        )

    left = decimal_fraction(BESSEL_LEFT)
    right = decimal_fraction(BESSEL_RIGHT)
    if right - left != Fraction(1, 10**49):
        raise AssertionError("The Bessel bracket width is not exactly 10^-49.")

    left_lower, left_upper, left_next = bessel_rational_enclosure(left)
    right_lower, right_upper, right_next = bessel_rational_enclosure(right)

    if not left_lower > 0:
        raise AssertionError("The exact lower bound at L is not positive.")
    if not right_upper < 0:
        raise AssertionError("The exact upper bound at U is not negative.")
    if not max(left_next, right_next) < Fraction(61, 10**126):
        raise AssertionError("The next-term bound 6.1e-125 failed.")

    print("\nExact-rational Bessel sign bracket")
    print(f"L = {BESSEL_LEFT}")
    print(f"U = {BESSEL_RIGHT}")
    print(f"U-L = 1e-49 exactly")
    print(f"J0(L) lower = {fraction_decimal(left_lower)}")
    print(f"J0(L) upper = {fraction_decimal(left_upper)}")
    print(f"J0(U) lower = {fraction_decimal(right_lower)}")
    print(f"J0(U) upper = {fraction_decimal(right_upper)}")
    print(
        "max next-term magnitude = "
        f"{fraction_decimal(max(left_next, right_next))}"
    )

    j1_lower = sum(
        Fraction((-1) ** k, factorial(k) * factorial(k + 1))
        for k in range(6)
    )
    if j1_lower != Fraction(49829, 86400):
        raise AssertionError("The J1(2) lower partial sum changed.")
    print(f"J1(2) lower partial sum = {j1_lower} > 0.5767")


def positive_real_square_root(value: mp.mpc) -> mp.mpc:
    """Choose sqrt(value) with nonnegative real part."""
    root = mp.sqrt(value)
    return root if mp.re(root) >= 0 else -root


def schwarzschild_coefficients(
    n: int,
    rho: mp.mpc,
) -> tuple[mp.mpc, mp.mpc, mp.mpc]:
    """Jaffé–Leaver coefficients for spin magnitude 2 and ell=2."""
    spin = mp.mpf(2)
    ell = 2
    angular = ell * (ell + 1)
    sigma = spin**2 - 1
    alpha = (n + 1) * (n + 2 * rho + 1)
    beta = -(
        2 * n**2
        + (8 * rho + 2) * n
        + 8 * rho**2
        + 4 * rho
        + angular
        - sigma
    )
    gamma = (n + 2 * rho) ** 2 - spin**2
    return alpha, beta, gamma


def schwarzschild_terminal_ratio(
    n: int,
    rho: mp.mpc,
    order: int,
) -> mp.mpc:
    """Large-order minimal-tail ratio R_n through order n^-2.

    order=0: zero tail
    order=1: 1-kappa*n^-1/2
    order=2: add n^-1
    order=3: add n^-3/2
    order=4: add n^-2
    """
    if order == 0:
        return mp.mpf("0")
    if order not in (1, 2, 3, 4):
        raise ValueError("Schwarzschild tail order must be 0 through 4.")

    spin = mp.mpf(2)
    ell = 2
    angular = ell * (ell + 1)
    kappa = positive_real_square_root(2 * rho)
    ratio = 1 - kappa / mp.sqrt(n)

    if order >= 2:
        ratio += (2 * rho - mp.mpf(3) / 4) / n

    c3 = (
        64 * rho**2
        - 48 * rho
        + 16 * angular
        + 3
    ) / (32 * kappa)
    if order >= 3:
        ratio -= c3 / n ** mp.mpf("1.5")

    c4 = (
        64 * rho * (angular - spin**2)
        + 48 * rho
        + 16 * angular
        + 3
    ) / (128 * rho)
    if order >= 4:
        ratio += c4 / n**2
    return ratio


def schwarzschild_residual(
    rho: mp.mpc,
    cutoff: int,
    tail_order: int,
) -> mp.mpc:
    """Evaluate the endpoint continued-fraction residual in rho=-i*omega."""
    ratio = schwarzschild_terminal_ratio(cutoff + 1, rho, tail_order)
    for n in range(cutoff, 0, -1):
        alpha, beta, gamma = schwarzschild_coefficients(n, rho)
        ratio = -gamma / (beta + alpha * ratio)
    alpha0, beta0, _ = schwarzschild_coefficients(0, rho)
    return beta0 + alpha0 * ratio


def schwarzschild_root(
    cutoff: int,
    tail_order: int,
    dps: int,
) -> mp.mpc:
    """Solve for the positive-real fundamental Schwarzschild QNM."""
    mp.mp.dps = dps
    omega0 = mp.mpc(QNM_INITIAL_OMEGA_RE, QNM_INITIAL_OMEGA_IM)
    rho0 = -mp.j * omega0
    rho1 = rho0 * (1 + mp.mpf("1e-6"))
    tolerance = mp.power(10, -(dps - 12))
    rho = mp.findroot(
        lambda value: schwarzschild_residual(value, cutoff, tail_order),
        (rho0, rho1),
        tol=tolerance,
        maxsteps=60,
    )
    return mp.j * rho


def print_qnm_audit() -> None:
    """Print the Schwarzschild cutoff, tail-order, and precision audits."""
    print("\nSchwarzschild recurrence audit")
    print("mode: gravitational spin magnitude s=2, ell=2, overtone n=0")
    print("convention: exp(-i*omega*t), 2M=1, rho=-i*omega")
    print(
        "initial omega: "
        f"{QNM_INITIAL_OMEGA_RE}{QNM_INITIAL_OMEGA_IM}j"
    )
    print(f"working decimal digits: {QNM_DPS}")
    print(f"tail order: {QNM_TAIL_ORDER}")

    roots = [
        schwarzschild_root(cutoff, QNM_TAIL_ORDER, QNM_DPS)
        for cutoff in QNM_CUTOFFS
    ]
    reference_400 = roots[-1]
    print("N       omega_N                                      |omega_N-omega_400|")
    for cutoff, root in zip(QNM_CUTOFFS, roots):
        difference = abs(root - reference_400)
        difference_text = (
            "reference row"
            if cutoff == QNM_CUTOFFS[-1]
            else mp.nstr(difference, 8)
        )
        print(
            f"{cutoff:4d}  "
            f"{mp.nstr(root, 40):>48}  "
            f"{difference_text}"
        )

    deep_reference = schwarzschild_root(
        QNM_REFERENCE_CUTOFF,
        QNM_TAIL_ORDER,
        QNM_DPS,
    )
    print(
        f"\nTail-order audit at N=200 against N={QNM_REFERENCE_CUTOFF}, K=4"
    )
    for order in range(5):
        root = schwarzschild_root(200, order, QNM_DPS)
        print(f"K={order}: {mp.nstr(abs(root-deep_reference), 10)}")

    print("\nPrecision audit at N=200, K=4")
    for dps in (30, 40, 50, 60, 80):
        root = schwarzschild_root(200, QNM_TAIL_ORDER, dps)
        print(f"p={dps:2d}: {mp.nstr(root, 25)}")


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--section",
        choices=("all", "bessel", "qnm"),
        default="all",
        help="Select which audit to run (default: all).",
    )
    arguments = parser.parse_args()

    print(f"CPython {platform.python_version()}")
    print(f"mpmath {mp.__version__}")
    if mp.__version__ != "1.3.0":
        print("warning: the published table used mpmath 1.3.0")

    if arguments.section in ("all", "bessel"):
        print_bessel_audit()
    if arguments.section in ("all", "qnm"):
        print_qnm_audit()


if __name__ == "__main__":
    main()