#!/usr/bin/env python3
"""Reproduce the sign and coefficient checks for Chapter 9, Page 2."""

from __future__ import annotations

import platform
import sys

import mpmath as mp
import sympy as sp


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

    if not condition:
        raise RuntimeError(message)


def require_close(
    actual: complex,
    expected: complex,
    tolerance: mp.mpf,
    message: str,
) -> None:
    """Compare complex values with a relative scale."""

    error = abs(actual - expected)
    scale = max(mp.mpf(1), abs(expected))
    if error > tolerance * scale:
        raise RuntimeError(
            f"{message}: error {mp.nstr(error, 8)} exceeds "
            f"{mp.nstr(tolerance * scale, 8)}"
        )


def main() -> int:
    mp.mp.dps = 70

    print("Advanced ODE Chapter 9 Page 2 audit")
    print(f"Python: {platform.python_version()}")
    print(f"SymPy: {sp.__version__}")
    print(f"mpmath: {mp.__version__}")

    xi = sp.symbols("xi")
    omega, beta = sp.symbols(
        "omega beta",
        positive=True,
    )

    # The shifted Borel transform of
    # c_n = Gamma(n+beta)/(Gamma(beta) omega**n)
    # is (1-xi/omega)**(-beta).
    branch = (1 - xi / omega) ** (-beta)
    branch_series = sp.series(branch, xi, 0, 8).removeO().expand()
    for n in range(8):
        formal_coefficient = (
            sp.gamma(n + beta)
            / (sp.gamma(beta) * omega**n)
        )
        borel_coefficient = (
            formal_coefficient / sp.factorial(n)
        )
        require(
            sp.simplify(
                branch_series.coeff(xi, n)
                - borel_coefficient
            )
            == 0,
            f"shifted Borel coefficient failed at n={n}",
        )
    print("PASS shifted Borel coefficients through n=7")

    # Page 1's upper-minus-lower contour is clockwise.
    # A residue r therefore contributes -2*pi*i*r.
    h = sp.symbols("h", positive=True)
    pole = 1 / (1 - xi / omega)
    weighted_residue = sp.residue(
        sp.exp(-xi / h) * pole,
        xi,
        omega,
    )
    pole_jump = sp.simplify(-2 * sp.pi * sp.I * weighted_residue)
    expected_pole_jump = (
        2 * sp.pi * sp.I * omega * sp.exp(-omega / h)
    )
    require(
        sp.simplify(pole_jump - expected_pole_jump) == 0,
        "clockwise pole jump has the wrong sign",
    )
    print("PASS pole jump sign: upper minus lower is clockwise")

    sample_omega = mp.mpf("1.3")
    sample_h = mp.mpf("0.27")
    sample_beta = mp.mpf("0.37")
    probe_s = mp.mpf("0.41")
    boundary_epsilon = mp.mpf("1e-50")

    # Evaluate the two principal-log boundary values rather than
    # inserting the desired jump sign into the Laplace integral.
    upper_xi = sample_omega + probe_s + mp.mpc(0, boundary_epsilon)
    lower_xi = sample_omega + probe_s - mp.mpc(0, boundary_epsilon)
    log_upper = -mp.log(1 - upper_xi / sample_omega)
    log_lower = -mp.log(1 - lower_xi / sample_omega)
    log_boundary_jump = log_upper - log_lower
    require_close(
        log_boundary_jump,
        2j * mp.pi,
        mp.mpf("1e-40"),
        "logarithmic upper-minus-lower boundary value",
    )

    log_jump_integral = (
        2j
        * mp.pi
        * mp.e ** (-sample_omega / sample_h)
        * mp.quad(
            lambda s: mp.e ** (-s / sample_h),
            [0, mp.inf],
        )
    )
    expected_log_jump = (
        2j
        * mp.pi
        * sample_h
        * mp.e ** (-sample_omega / sample_h)
    )
    require_close(
        log_jump_integral,
        expected_log_jump,
        mp.mpf("1e-60"),
        "logarithmic jump",
    )
    print(
        "PASS logarithmic jump: "
        f"{mp.nstr(log_jump_integral, 18)}"
    )

    # For 0 < beta < 1, independently evaluate the algebraic branch
    # boundary factor and then integrate it along the cut.
    branch_upper = (1 - upper_xi / sample_omega) ** (-sample_beta)
    branch_lower = (1 - lower_xi / sample_omega) ** (-sample_beta)
    branch_boundary_jump = branch_upper - branch_lower
    expected_branch_boundary_jump = (
        2j
        * mp.sin(mp.pi * sample_beta)
        * (probe_s / sample_omega) ** (-sample_beta)
    )
    require_close(
        branch_boundary_jump,
        expected_branch_boundary_jump,
        mp.mpf("1e-40"),
        "algebraic upper-minus-lower boundary value",
    )

    branch_jump_integral = (
        2j
        * mp.sin(mp.pi * sample_beta)
        * sample_omega**sample_beta
        * mp.e ** (-sample_omega / sample_h)
        * mp.quad(
            lambda s: (
                mp.e ** (-s / sample_h)
                * s ** (-sample_beta)
            ),
            [0, 1, mp.inf],
        )
    )
    expected_branch_jump = (
        2j
        * mp.pi
        * sample_omega**sample_beta
        * sample_h ** (1 - sample_beta)
        * mp.e ** (-sample_omega / sample_h)
        / mp.gamma(sample_beta)
    )
    require_close(
        branch_jump_integral,
        expected_branch_jump,
        mp.mpf("1e-45"),
        "algebraic branch jump",
    )
    print(
        "PASS algebraic branch jump: "
        f"{mp.nstr(branch_jump_integral, 18)}"
    )

    reflection_left = (
        mp.sin(mp.pi * sample_beta)
        * mp.gamma(1 - sample_beta)
    )
    reflection_right = mp.pi / mp.gamma(sample_beta)
    require_close(
        reflection_left,
        reflection_right,
        mp.mpf("1e-65"),
        "gamma reflection identity",
    )
    print("PASS gamma reflection normalization")

    # Exact ratios recover both the singularity location and exponent.
    n = 35

    def coefficient(index: int) -> mp.mpf:
        return (
            mp.gamma(index + sample_beta)
            / (
                mp.gamma(sample_beta)
                * sample_omega**index
            )
        )

    ratio_n = coefficient(n + 1) / coefficient(n)
    ratio_previous = coefficient(n) / coefficient(n - 1)
    recovered_omega = 1 / (ratio_n - ratio_previous)
    recovered_beta = recovered_omega * ratio_n - n
    require_close(
        recovered_omega,
        sample_omega,
        mp.mpf("1e-60"),
        "large-order recovery of omega",
    )
    require_close(
        recovered_beta,
        sample_beta,
        mp.mpf("1e-60"),
        "large-order recovery of beta",
    )
    print(
        "PASS large-order recovery: "
        f"omega={mp.nstr(recovered_omega, 16)}, "
        f"beta={mp.nstr(recovered_beta, 16)}"
    )

    # Verify the two-pole convolution near the base germ.
    # Rational sample values avoid assumptions about symbolic branches.
    eta = sp.symbols("eta")
    aval = sp.Rational(2)
    bval = sp.Rational(3)
    integrand = 1 / (
        (aval - eta)
        * (bval - xi + eta)
    )
    convolution_integral = sp.integrate(
        integrand,
        (eta, 0, xi),
    )
    convolution_closed = (
        sp.log(aval / (aval - xi))
        + sp.log(bval / (bval - xi))
    ) / (aval + bval - xi)
    integral_series = sp.series(
        convolution_integral,
        xi,
        0,
        9,
    ).removeO()
    closed_series = sp.series(
        convolution_closed,
        xi,
        0,
        9,
    ).removeO()
    require(
        sp.simplify(integral_series - closed_series) == 0,
        "two-pole convolution identity failed",
    )

    later_sheet_term = (
        2 * sp.pi * sp.I / (aval + bval - xi)
    )
    later_sheet_residue = sp.residue(
        later_sheet_term,
        xi,
        aval + bval,
    )
    require(
        sp.simplify(later_sheet_residue + 2 * sp.pi * sp.I)
        == 0,
        "later-sheet residue failed",
    )
    print("PASS convolution identity and later-sheet pole")

    # Check the action-2A grade of exp(t*Delta_A+t**2*Delta_2A)
    # with genuinely noncommuting matrices.
    t = sp.symbols("t")
    delta_one = sp.Matrix([[0, 1], [1, 0]])
    delta_two = sp.Matrix([[1, 0], [0, -1]])
    identity = sp.eye(2)
    generator = t * delta_one + t**2 * delta_two
    exponential_to_grade_two = (
        identity
        + generator
        + generator * generator / 2
    )
    expected_to_grade_two = (
        identity
        + t * delta_one
        + t**2 * (
            delta_two + delta_one * delta_one / 2
        )
    )
    for row in range(2):
        for column in range(2):
            difference = sp.series(
                exponential_to_grade_two[row, column]
                - expected_to_grade_two[row, column],
                t,
                0,
                3,
            ).removeO()
            require(
                sp.expand(difference) == 0,
                "Stokes exponential grading failed",
            )
    print("PASS Stokes automorphism through action grade 2A")

    # The Stirling Borel transform has residue 1/(2*pi*i*m).
    stirling_borel = (
        (xi / 2) * sp.coth(xi / 2) - 1
    ) / xi**2
    for m in (1, 2):
        singularity = 2 * sp.pi * sp.I * m
        residue = sp.simplify(
            sp.limit(
                (xi - singularity) * stirling_borel,
                xi,
                singularity,
            )
        )
        expected_residue = 1 / (2 * sp.pi * sp.I * m)
        require(
            sp.simplify(residue - expected_residue) == 0,
            f"Stirling residue failed at m={m}",
        )
    print("PASS first two Stirling-lattice residues")

    print("All resurgent-singularity checks passed.")
    return 0


if __name__ == "__main__":
    try:
        raise SystemExit(main())
    except Exception as error:
        print(f"FAIL {error}", file=sys.stderr)
        raise SystemExit(1)