#!/usr/bin/env python3
"""Optimization-safe audits for Chapter 9, Page 9.

The script keeps two calibrations deliberately separate.

1. The regularized Weber Voros exponent tests shifted Borel coefficients,
   parity-adapted Padé continuation, the exact imaginary pole lattice,
   a nonsingular Borel--Laplace sum, and a lateral Stokes jump.
2. The PT-symmetric Airy box tests determinant normalization, arbitrary-
   precision Taylor shooting, complex root polishing, and an argument-
   principle root count.

Every failure raises an explicit exception.  No Python assert statements are
used, so normal and optimized execution perform the same audit.  The checks
provide strong cross-verification, but they do not turn finite Padé data or a
sampled winding number into interval-certified analytic continuation.
"""

from __future__ import annotations

import platform
from functools import lru_cache
from typing import Callable, Iterable, Sequence, Tuple

import mpmath as mp
import sympy as sp


def require(condition: bool, message: str) -> None:
    """Raise an optimization-safe failure unless ``condition`` is true."""

    if not condition:
        raise RuntimeError(message)


def require_close(
    actual: complex,
    expected: complex,
    tolerance: float,
    label: str,
) -> None:
    """Require an absolute numerical agreement at a declared tolerance."""

    error = abs(actual - expected)
    if not error < tolerance:
        raise RuntimeError(
            f"{label}: error {mp.nstr(error, 12)} exceeds "
            f"{mp.nstr(tolerance, 12)}"
        )


def polynomial_value(
    coefficients: Sequence[complex],
    argument: complex,
) -> complex:
    """Evaluate an ascending-order coefficient list by Horner's rule."""

    value = mp.mpc(0)
    for coefficient in reversed(coefficients):
        value = value * argument + coefficient
    return value


def pade_from_series(
    coefficients: Sequence[complex],
    numerator_degree: int,
    denominator_degree: int,
) -> Tuple[list, list]:
    """Construct [L/M] with Q(0)=1 from c_0 through c_{L+M}."""

    degree_sum = numerator_degree + denominator_degree
    require(
        len(coefficients) >= degree_sum + 1,
        "insufficient coefficients for the requested Padé shape",
    )
    if denominator_degree == 0:
        return list(coefficients[: numerator_degree + 1]), [mp.mpf(1)]

    matrix = mp.matrix(denominator_degree)
    right_hand_side = mp.matrix(denominator_degree, 1)
    for row in range(denominator_degree):
        degree = numerator_degree + 1 + row
        right_hand_side[row] = -coefficients[degree]
        for column in range(denominator_degree):
            denominator_index = column + 1
            series_index = degree - denominator_index
            matrix[row, column] = (
                coefficients[series_index]
                if series_index >= 0
                else mp.mpf(0)
            )

    solution = mp.lu_solve(matrix, right_hand_side)
    denominator = [mp.mpf(1)] + [
        solution[index] for index in range(denominator_degree)
    ]
    numerator = []
    for degree in range(numerator_degree + 1):
        numerator.append(
            sum(
                denominator[index] * coefficients[degree - index]
                for index in range(min(degree, denominator_degree) + 1)
            )
        )
    return numerator, denominator


def pade_matching_audit(
    coefficients: Sequence[complex],
    numerator: Sequence[complex],
    denominator: Sequence[complex],
) -> None:
    """Check Q times the input series minus P through the matched order."""

    matched_degree = len(numerator) + len(denominator) - 2
    for degree in range(matched_degree + 1):
        convolution = sum(
            denominator[index] * coefficients[degree - index]
            for index in range(min(degree, len(denominator) - 1) + 1)
        )
        numerator_coefficient = (
            numerator[degree] if degree < len(numerator) else mp.mpf(0)
        )
        require_close(
            convolution,
            numerator_coefficient,
            mp.mpf("1e-65"),
            f"Padé coefficient matching at degree {degree}",
        )


def sympy_rational_to_mpf(value: sp.Expr) -> mp.mpf:
    """Convert an exact SymPy rational to the current mpmath precision."""

    rational = sp.Rational(value)
    return mp.mpf(int(rational.p)) / int(rational.q)


def weber_u_borel_coefficients(count: int, mu: complex = 1) -> list:
    """Return b_j in B W(xi)=sum b_j (xi^2)^j for the Weber path."""

    mu_value = mp.mpf(mu)
    coefficients = []
    for index in range(count):
        k = index + 1
        bernoulli = sympy_rational_to_mpf(sp.bernoulli(2 * k))
        formal_coefficient = (
            (mp.power(2, 1 - 2 * k) - 1)
            * bernoulli
            / (2 * k * (2 * k - 1) * mu_value ** (2 * k - 1))
        )
        coefficients.append(formal_coefficient / mp.factorial(2 * k - 2))
    return coefficients


def weber_u_borel_coefficients_exact(count: int) -> list:
    """Return the mu=1 Weber coefficients as exact SymPy rationals."""

    coefficients = []
    for index in range(count):
        k = index + 1
        coefficients.append(
            (sp.Rational(2) ** (1 - 2 * k) - 1)
            * sp.bernoulli(2 * k)
            / sp.factorial(2 * k)
        )
    return coefficients


def sympy_pade_from_series(
    coefficients: Sequence[sp.Expr],
    numerator_degree: int,
    denominator_degree: int,
) -> Tuple[tuple, tuple]:
    """Construct an exact rational Padé approximant with Q(0)=1."""

    degree_sum = numerator_degree + denominator_degree
    require(
        len(coefficients) >= degree_sum + 1,
        "insufficient exact coefficients for the requested Padé shape",
    )
    if denominator_degree == 0:
        return tuple(coefficients[: numerator_degree + 1]), (sp.Integer(1),)

    def coefficient_at(index: int) -> sp.Expr:
        """Return a series coefficient, with the formal convention c_j=0 for j<0."""

        return coefficients[index] if 0 <= index < len(coefficients) else sp.Integer(0)

    matrix = sp.Matrix(
        denominator_degree,
        denominator_degree,
        lambda row, column: coefficient_at(
            numerator_degree + 1 + row - (column + 1)
        ),
    )
    right_hand_side = sp.Matrix(
        [
            -coefficients[numerator_degree + 1 + row]
            for row in range(denominator_degree)
        ]
    )
    solution = matrix.inv() * right_hand_side
    denominator = (sp.Integer(1),) + tuple(solution)
    numerator = tuple(
        sp.factor(
            sum(
                denominator[index] * coefficients[degree - index]
                for index in range(min(degree, denominator_degree) + 1)
            )
        )
        for degree in range(numerator_degree + 1)
    )

    matched_degree = numerator_degree + denominator_degree
    for degree in range(matched_degree + 1):
        convolution = sum(
            denominator[index] * coefficients[degree - index]
            for index in range(min(degree, denominator_degree) + 1)
        )
        target = numerator[degree] if degree <= numerator_degree else 0
        require(
            sp.factor(convolution - target) == 0,
            f"exact Weber Padé mismatch at degree {degree}",
        )
    return numerator, denominator


@lru_cache(maxsize=None)
def exact_weber_pade(count: int) -> Tuple[tuple, tuple, tuple]:
    """Cache one exact parity-adapted Weber Padé approximant."""

    coefficients = tuple(weber_u_borel_coefficients_exact(count))
    numerator_degree = count // 2
    denominator_degree = count - 1 - numerator_degree
    numerator, denominator = sympy_pade_from_series(
        coefficients,
        numerator_degree,
        denominator_degree,
    )
    return coefficients, numerator, denominator


def weber_pade(count: int) -> Tuple[list, list, list]:
    """Build a near-diagonal rational continuation in u=xi^2."""

    exact_coefficients, exact_numerator, exact_denominator = exact_weber_pade(
        count
    )
    coefficients = [sympy_rational_to_mpf(value) for value in exact_coefficients]
    numerator = [sympy_rational_to_mpf(value) for value in exact_numerator]
    denominator = [sympy_rational_to_mpf(value) for value in exact_denominator]
    return coefficients, numerator, denominator


def weber_pade_value(
    xi: complex,
    numerator: Sequence[complex],
    denominator: Sequence[complex],
) -> complex:
    """Evaluate the parity-adapted Weber Padé approximant."""

    argument = xi * xi
    return polynomial_value(numerator, argument) / polynomial_value(
        denominator,
        argument,
    )


def weber_positive_sum(count: int, hbar: complex) -> complex:
    """Evaluate the shifted positive-real Padé--Laplace integral."""

    _, numerator, denominator = weber_pade(count)
    hbar_value = mp.mpf(hbar)
    integrand = lambda variable: (
        mp.exp(-variable)
        * weber_pade_value(
            hbar_value * variable,
            numerator,
            denominator,
        )
    )
    return hbar_value * mp.quad(integrand, [0, 1, mp.inf])


def exact_weber_positive_sum(hbar: complex, mu: complex = 1) -> complex:
    """Return the exact gamma-function calibration of the Weber sum."""

    t_value = mp.mpf(mu) / mp.mpf(hbar)
    return (
        t_value
        + mp.loggamma(t_value + mp.mpf("0.5"))
        - mp.log(2 * mp.pi) / 2
        - t_value * mp.log(t_value)
    )


def weber_lateral_sum(count: int, angle: complex) -> complex:
    """Integrate the Weber Padé germ along a rotated ray for hbar=i."""

    _, numerator, denominator = weber_pade(count)
    rotation = mp.exp(mp.j * angle)

    def integrand(radius: complex) -> complex:
        xi = rotation * radius
        return (
            rotation
            * mp.exp(-xi / mp.j)
            * weber_pade_value(xi, numerator, denominator)
        )

    breakpoints = [
        0,
        2 * mp.pi,
        4 * mp.pi,
        6 * mp.pi,
        8 * mp.pi,
        12 * mp.pi,
        mp.inf,
    ]
    return mp.quad(integrand, breakpoints)


def weber_lateral_jump(count: int, offset: complex) -> complex:
    """Return S_{pi/2+delta} minus S_{pi/2-delta} at hbar=i."""

    central_angle = mp.pi / 2
    return weber_lateral_sum(
        count,
        central_angle + offset,
    ) - weber_lateral_sum(
        count,
        central_angle - offset,
    )


def weber_pole_audit() -> None:
    """Compare the first Padé poles with xi_k=2*pi*i*k."""

    _, _, denominator = exact_weber_pade(16)
    variable = sp.symbols("u")
    polynomial = sum(
        coefficient * variable**degree
        for degree, coefficient in enumerate(denominator)
    )
    roots_u = sp.nroots(polynomial, n=50, maxsteps=300)
    roots_xi = []
    for root in roots_u:
        root_mpc = mp.mpc(
            str(sp.re(root).evalf(55)),
            str(sp.im(root).evalf(55)),
        )
        principal_root = mp.sqrt(root_mpc)
        roots_xi.extend((principal_root, -principal_root))
    positive_poles = [root for root in roots_xi if mp.im(root) > 0]
    for index in range(1, 4):
        expected = 2 * mp.pi * mp.j * index
        closest = min(positive_poles, key=lambda root: abs(root - expected))
        require(
            abs(closest - expected) < mp.mpf("2e-6"),
            "Weber Padé pole k={}: error {}".format(
                index,
                mp.nstr(abs(closest - expected), 12),
            ),
        )
        print(
            "  Weber pole k={}: {}".format(
                index,
                mp.nstr(closest, 18),
            )
        )


def weber_audit() -> None:
    """Run nonsingular, pole-lattice, and lateral Weber calibrations."""

    mp.mp.dps = 80
    sample_coefficients, numerator, denominator = weber_pade(16)
    pade_matching_audit(sample_coefficients, numerator, denominator)

    # Exercise the c_j=0 convention needed by subdiagonal shapes.  The
    # [0/2] approximant would be corrupted by Python negative indexing if
    # the missing c_{-1} entry were not inserted explicitly.
    test_numerator, test_denominator = pade_from_series(
        [mp.mpf(1), mp.mpf(1), mp.mpf("0.5")],
        0,
        2,
    )
    pade_matching_audit(
        [mp.mpf(1), mp.mpf(1), mp.mpf("0.5")],
        test_numerator,
        test_denominator,
    )

    # Check the same convention in exact arithmetic on a near-diagonal
    # Weber shape with M=L+2.
    subdiagonal_coefficients = weber_u_borel_coefficients_exact(11)
    sympy_pade_from_series(subdiagonal_coefficients, 4, 6)

    hbar = mp.mpf("0.2")
    exact_sum = exact_weber_positive_sum(hbar)
    errors = []
    for count in (4, 8, 12, 20):
        approximation = weber_positive_sum(count, hbar)
        errors.append(abs(approximation - exact_sum))
        print(
            "  Weber S_0 K={:2d} [{}/{}]: error {}".format(
                count,
                count // 2,
                count - 1 - count // 2,
                mp.nstr(errors[-1], 8),
            )
        )
    require(
        all(errors[index + 1] < errors[index] for index in range(3)),
        "Weber nonsingular Borel--Padé errors did not decrease",
    )
    require(
        errors[-1] < mp.mpf("2e-48"),
        "Weber nonsingular Borel--Padé target was not reached",
    )

    weber_pole_audit()

    exact_jump = mp.log1p(mp.exp(-2 * mp.pi))
    jump_errors = []
    for count in (8, 16):
        jump = weber_lateral_jump(count, mp.mpf("0.1"))
        require_close(
            mp.im(jump),
            0,
            mp.mpf("1e-65"),
            f"Weber lateral jump reality at K={count}",
        )
        jump_errors.append(abs(jump - exact_jump))
        print(
            "  Weber jump K={:2d} [{}/{}]: error {}".format(
                count,
                count // 2,
                count - 1 - count // 2,
                mp.nstr(jump_errors[-1], 8),
            )
        )
    require(
        jump_errors[1] < mp.mpf("1e-17"),
        "Weber lateral jump target was not reached",
    )

    offset_check = weber_lateral_jump(16, mp.mpf("0.12"))
    require_close(
        offset_check,
        exact_jump,
        mp.mpf("1e-17"),
        "Weber lateral contour-offset stability",
    )


def airy_delta(energy: complex, coupling: complex) -> complex:
    """Return the PT-real-normalized Airy-box determinant Delta_A."""

    coupling_value = mp.mpf(coupling)
    alpha = coupling_value ** (mp.mpf(1) / 3) * mp.exp(mp.j * mp.pi / 6)
    z_minus = alpha * (-1 + mp.j * energy / coupling_value)
    z_plus = alpha * (1 + mp.j * energy / coupling_value)
    raw = (
        mp.airyai(z_minus) * mp.airybi(z_plus)
        - mp.airyai(z_plus) * mp.airybi(z_minus)
    )
    return mp.exp(-mp.j * mp.pi / 6) * raw


def airy_delta_derivative(energy: complex, coupling: complex) -> complex:
    """Differentiate the PT-real-normalized Airy determinant in energy."""

    coupling_value = mp.mpf(coupling)
    alpha = coupling_value ** (mp.mpf(1) / 3) * mp.exp(mp.j * mp.pi / 6)
    z_minus = alpha * (-1 + mp.j * energy / coupling_value)
    z_plus = alpha * (1 + mp.j * energy / coupling_value)
    dz_de = alpha * mp.j / coupling_value
    derivative_raw = dz_de * (
        mp.airyai(z_minus, 1) * mp.airybi(z_plus)
        + mp.airyai(z_minus) * mp.airybi(z_plus, 1)
        - mp.airyai(z_plus, 1) * mp.airybi(z_minus)
        - mp.airyai(z_plus) * mp.airybi(z_minus, 1)
    )
    return mp.exp(-mp.j * mp.pi / 6) * derivative_raw


def taylor_step(
    x_start: complex,
    value: complex,
    derivative: complex,
    step: complex,
    energy: complex,
    coupling: complex,
    order: int,
) -> Tuple[complex, complex]:
    """Advance y''=(i*g*x-E)y by one arbitrary-precision Taylor step."""

    coefficients = [mp.mpc(0) for _ in range(order + 1)]
    coefficients[0] = value
    coefficients[1] = derivative
    for index in range(order - 1):
        previous = coefficients[index - 1] if index > 0 else mp.mpc(0)
        coefficients[index + 2] = (
            (mp.j * coupling * x_start - energy) * coefficients[index]
            + mp.j * coupling * previous
        ) / ((index + 2) * (index + 1))

    advanced_value = polynomial_value(coefficients, step)
    derivative_coefficients = [
        index * coefficients[index] for index in range(1, order + 1)
    ]
    advanced_derivative = polynomial_value(derivative_coefficients, step)
    return advanced_value, advanced_derivative


def airy_shooting_determinant(
    energy: complex,
    coupling: complex,
    steps: int = 8,
    order: int = 60,
) -> complex:
    """Shoot from y(-1)=0,y'(-1)=1 and return y(1)."""

    x_value = mp.mpf(-1)
    value = mp.mpc(0)
    derivative = mp.mpc(1)
    step = mp.mpf(2) / steps
    for _ in range(steps):
        value, derivative = taylor_step(
            x_value,
            value,
            derivative,
            step,
            energy,
            coupling,
            order,
        )
        x_value += step
    return value


def airy_exact_root(guess: complex, coupling: complex) -> complex:
    """Polish one exact Airy-determinant zero."""

    start = mp.mpc(guess)
    return mp.findroot(
        lambda variable: airy_delta(variable, coupling),
        (start, start + mp.mpf("1e-4")),
        tol=mp.mpf("1e-60"),
        maxsteps=80,
    )


def airy_shooting_root(
    guess: complex,
    coupling: complex,
    steps: int = 8,
    order: int = 60,
) -> complex:
    """Polish one zero using only the Taylor-shooting determinant."""

    start = mp.mpc(guess)
    return mp.findroot(
        lambda variable: airy_shooting_determinant(
            variable,
            coupling,
            steps=steps,
            order=order,
        ),
        (start, start + mp.mpf("1e-4")),
        tol=mp.mpf("1e-58"),
        maxsteps=80,
    )


def rectangle_points(
    x_min: float,
    x_max: float,
    y_min: float,
    y_max: float,
    points_per_edge: int,
) -> list:
    """Return a counterclockwise sample of a rectangle without duplicates."""

    points = []
    for index in range(points_per_edge):
        parameter = mp.mpf(index) / points_per_edge
        points.append(mp.mpc(x_min + (x_max - x_min) * parameter, y_min))
    for index in range(points_per_edge):
        parameter = mp.mpf(index) / points_per_edge
        points.append(mp.mpc(x_max, y_min + (y_max - y_min) * parameter))
    for index in range(points_per_edge):
        parameter = mp.mpf(index) / points_per_edge
        points.append(mp.mpc(x_max - (x_max - x_min) * parameter, y_max))
    for index in range(points_per_edge):
        parameter = mp.mpf(index) / points_per_edge
        points.append(mp.mpc(x_min, y_max - (y_max - y_min) * parameter))
    return points


def sampled_winding(
    function: Callable[[complex], complex],
    contour: Sequence[complex],
) -> Tuple[int, complex, complex]:
    """Compute a phase-unwrapped sampled winding and boundary clearance."""

    values = [function(point) for point in contour]
    require(
        min(abs(value) for value in values) > mp.mpf("1e-5"),
        "argument-principle contour passes too close to a sampled zero",
    )
    phase_increments = []
    for index, value in enumerate(values):
        next_value = values[(index + 1) % len(values)]
        phase_increments.append(mp.arg(next_value / value))
    maximum_increment = max(abs(increment) for increment in phase_increments)
    require(
        maximum_increment < mp.pi / 2,
        "sampled Airy winding has an aliased phase increment",
    )
    phase_change = sum(phase_increments)
    winding_float = phase_change / (2 * mp.pi)
    winding = int(mp.nint(winding_float))
    require_close(
        winding_float,
        winding,
        mp.mpf("2e-12"),
        "sampled Airy argument-principle winding",
    )
    return winding, min(abs(value) for value in values), maximum_increment


def rectangle_log_derivative_winding(
    function: Callable[[complex], complex],
    derivative: Callable[[complex], complex],
    x_min: float,
    x_max: float,
    y_min: float,
    y_max: float,
) -> complex:
    """Integrate D_E/D independently along the four rectangle edges."""

    corners = (
        mp.mpc(x_min, y_min),
        mp.mpc(x_max, y_min),
        mp.mpc(x_max, y_max),
        mp.mpc(x_min, y_max),
    )
    integral = mp.mpc(0)
    for start, end in zip(corners, corners[1:] + corners[:1]):
        displacement = end - start
        integral += mp.quad(
            lambda parameter: (
                derivative(start + displacement * parameter)
                / function(start + displacement * parameter)
                * displacement
            ),
            [0, 1],
        )
    return integral / (2 * mp.pi * mp.j)


def airy_audit() -> None:
    """Cross-check the Airy determinant, roots, and root count."""

    mp.mp.dps = 75
    coupling = mp.mpf(10)
    normalization = mp.pi / coupling ** (mp.mpf(1) / 3)
    samples: Iterable[complex] = (
        mp.mpf("4.2"),
        mp.mpf("7.1"),
        mp.mpc("6.3", "0.4"),
    )
    for energy in samples:
        shooting_value = airy_shooting_determinant(energy, coupling)
        exact_value = normalization * airy_delta(energy, coupling)
        require_close(
            shooting_value,
            exact_value,
            mp.mpf("2e-65"),
            f"Airy Taylor-shooting normalization at E={energy}",
        )

    expected_roots = (
        mp.mpf("4.578490418291691965603793827641933036460687180596"),
        mp.mpf("8.996683255712348208437330264006417137735072636008"),
    )
    shooting_roots = []
    for guess, expected in zip((4.6, 9.0), expected_roots):
        exact_root = airy_exact_root(guess, coupling)
        shooting_root = airy_shooting_root(guess, coupling)
        alternate_shooting_root = airy_shooting_root(
            guess,
            coupling,
            steps=10,
            order=56,
        )
        require_close(
            exact_root,
            expected,
            mp.mpf("2e-48"),
            "exact Airy-box root",
        )
        require_close(
            shooting_root,
            exact_root,
            mp.mpf("2e-58"),
            "independent Taylor-shooting Airy root",
        )
        require_close(
            alternate_shooting_root,
            exact_root,
            mp.mpf("2e-54"),
            "refined-step Taylor-shooting Airy root",
        )
        shooting_roots.append(shooting_root)

    for index, root in enumerate(shooting_roots):
        print(f"  Airy E_{index}: {mp.nstr(root, 52)}")

    contour_80 = rectangle_points(3, 10, -1, 1, 80)
    winding_80, clearance_80, increment_80 = sampled_winding(
        lambda energy: airy_delta(energy, coupling),
        contour_80,
    )
    contour_160 = rectangle_points(3, 10, -1, 1, 160)
    winding, clearance, increment_160 = sampled_winding(
        lambda energy: airy_delta(energy, coupling),
        contour_160,
    )
    require(
        winding_80 == winding,
        "Airy sampled winding changed under contour-node doubling",
    )
    log_derivative_winding = rectangle_log_derivative_winding(
        lambda energy: airy_delta(energy, coupling),
        lambda energy: airy_delta_derivative(energy, coupling),
        3,
        10,
        -1,
        1,
    )
    require_close(
        log_derivative_winding,
        winding,
        mp.mpf("2e-45"),
        "Airy log-derivative winding",
    )
    require(winding == 2, f"Airy rectangle contains {winding} roots, not 2")
    print(f"  Airy rectangle winding (80/160 nodes and D_E/D): {winding}")
    print(f"  Sampled determinant clearance: {mp.nstr(clearance, 8)}")
    print(
        "  Maximum phase increment (80/160): {}/{}".format(
            mp.nstr(increment_80, 8),
            mp.nstr(increment_160, 8),
        )
    )
    require(
        clearance_80 > mp.mpf("1e-5"),
        "coarse Airy contour has inadequate sampled clearance",
    )


def main() -> None:
    """Run every deterministic Page 9 audit."""

    print("Advanced ODE Chapter 9 Page 9 Borel--Pade spectral audit")
    print(f"Python: {platform.python_version()}")
    print(f"SymPy: {sp.__version__}")
    print(f"mpmath: {mp.__version__}")
    print("Weber settings: mu=1, hbar=0.2 or i, shifted Borel transform")
    print("Airy settings: g=10, 75 dps, Taylor grids (8,60) and (10,56)")

    weber_audit()
    print("PASS Weber Padé matching, pole lattice, gamma sum, and lateral jump")

    airy_audit()
    print("PASS Airy determinant, two Taylor grids, roots, and contour counts")

    print(
        "LIMITATION: Padé-order stability and contour counts are numerical "
        "evidence, not interval-certified continuation or root enclosure."
    )
    print("All Borel--Padé and spectral checks passed.")


if __name__ == "__main__":
    main()