#!/usr/bin/env python3
"""Exact audit of the regularized WKB period formulas on Chapter 8, Page 5.

The book uses

    hbar**2 psi'' = R psi,
    P**2 + hbar P' = R,
    Omega = P_even dz = sum_{k >= 0} hbar**(2*k) lambda_(2*k).

This script checks four calculations that make the endpoint and coordinate
conventions concrete:

1. Airy turning-point normalization by half of a lifted contour;
2. the exact Riccati residue at a double pole and its Langer cancellation;
3. the Weber closed period and the first three regularized relative-period
   coefficients; and
4. cancellation of a spurious order-hbar**2 term by the Schwarzian under
   z = exp(w).

All calculations are exact.  Every failure raises ``RuntimeError`` explicitly,
so Python's ``-O`` flag does not disable the audit.  The script checks formal
coefficients only; it makes no claim about Borel summation or exact
quantization.

Compatibility target: Python 3.9 and SymPy 1.14.
"""

from __future__ import annotations

import platform

import sympy as sp


def reduced(expression: sp.Expr) -> sp.Expr:
    """Return a stable exact form for a rational symbolic expression."""
    return sp.factor(sp.cancel(sp.together(expression)))


def require_zero(expression: sp.Expr, label: str) -> None:
    """Raise a visible failure unless an exact expression vanishes."""
    residual = reduced(expression)
    if residual != 0:
        residual = sp.simplify(residual)
    if residual != 0:
        raise RuntimeError(f"{label} failed: residual = {residual}")


def generate_branch_on_cover(
    leading_momentum: sp.Expr,
    base_coordinate: sp.Expr,
    cover_coordinate: sp.Symbol,
    order: int,
) -> list[sp.Expr]:
    """Generate one Riccati branch after rationalizing the spectral cover."""
    if order < 0:
        raise RuntimeError(f"formal order must be nonnegative, got {order}")

    dz_du = sp.diff(base_coordinate, cover_coordinate)
    if dz_du == 0:
        raise RuntimeError("the cover parametrization has zero derivative")

    def differentiate_in_z(expression: sp.Expr) -> sp.Expr:
        return reduced(sp.diff(expression, cover_coordinate) / dz_du)

    coefficients = [reduced(leading_momentum)]
    for n in range(1, order + 1):
        convolution = sum(
            coefficients[j] * coefficients[n - j]
            for j in range(1, n)
        )
        next_coefficient = -(
            differentiate_in_z(coefficients[n - 1]) + convolution
        ) / (2 * coefficients[0])
        coefficients.append(reduced(next_coefficient))

    # Recheck the recurrence independently of the assignment above.
    for n in range(1, order + 1):
        convolution = sum(
            coefficients[j] * coefficients[n - j]
            for j in range(1, n)
        )
        residual = (
            2 * coefficients[0] * coefficients[n]
            + convolution
            + differentiate_in_z(coefficients[n - 1])
        )
        require_zero(residual, f"Riccati recursion coefficient hbar^{n}")

    return coefficients


def local_residue_at_infinity(
    one_form_coefficient: sp.Expr,
    coordinate: sp.Symbol,
) -> sp.Expr:
    """Return the residue using the positive local coordinate u=1/coordinate."""
    u = sp.symbols(f"u_{coordinate}")
    pulled_coefficient = -one_form_coefficient.subs(coordinate, 1 / u) / u**2
    return reduced(sp.residue(pulled_coefficient, u, 0))


def airy_turning_point_audit() -> dict[str, sp.Expr]:
    """Check the Airy order-hbar**2 half-contour regularization."""
    xi = sp.symbols("xi", nonzero=True)
    z = xi**2
    y = xi
    dz_dxi = sp.diff(z, xi)
    coefficients = generate_branch_on_cover(y, z, xi, 2)

    lambda_zero = reduced(coefficients[0] * dz_dxi)
    lambda_two = reduced(coefficients[2] * dz_dxi)
    expected_lambda_two = -sp.Rational(5, 16) / xi**4
    require_zero(
        lambda_two - expected_lambda_two,
        "Airy lambda_2 in the ramification coordinate",
    )
    require_zero(
        sp.residue(lambda_two, xi, 0),
        "Airy lambda_2 has no turning-point residue",
    )

    primitive_zero = sp.Rational(2, 3) * xi**3
    primitive_two = sp.Rational(5, 48) / xi**3
    require_zero(
        sp.diff(primitive_zero, xi) - lambda_zero,
        "Airy classical primitive",
    )
    require_zero(
        sp.diff(primitive_two, xi) - lambda_two,
        "Airy lambda_2 primitive",
    )

    half_classical = reduced(
        (primitive_zero - primitive_zero.subs(xi, -xi)) / 2
    )
    half_quantum = reduced(
        (primitive_two - primitive_two.subs(xi, -xi)) / 2
    )
    require_zero(
        half_classical - sp.Rational(2, 3) * xi**3,
        "Airy classical half-contour action",
    )
    require_zero(
        half_quantum - sp.Rational(5, 48) / xi**3,
        "Airy order-hbar**2 half-contour action",
    )

    return {
        "lambda_two": lambda_two,
        "half_classical": half_classical,
        "half_quantum": half_quantum,
    }


def double_pole_audit() -> dict[str, sp.Expr]:
    """Check the exact branch-difference residue and the Langer shift."""
    t = sp.symbols("t")
    c = sp.symbols("c", positive=True)
    hbar = sp.symbols("hbar")

    # Source-free inverse-square model R=c**2/t**2.
    coefficients = generate_branch_on_cover(c / t, t, t, 4)
    residue_two = reduced(sp.residue(coefficients[2], t, 0))
    residue_four = reduced(sp.residue(coefficients[4], t, 0))
    require_zero(
        residue_two - 1 / (8 * c),
        "source-free double-pole residue of lambda_2",
    )
    require_zero(
        residue_four + 1 / (128 * c**3),
        "source-free double-pole residue of lambda_4",
    )

    exact_residue = sp.sqrt(c**2 + hbar**2 / 4)
    source_free_plus = (hbar + sp.sqrt(hbar**2 + 4 * c**2)) / 2
    source_free_minus = (hbar - sp.sqrt(hbar**2 + 4 * c**2)) / 2
    require_zero(
        source_free_plus**2 - hbar * source_free_plus - c**2,
        "source-free plus Riccati residue",
    )
    require_zero(
        source_free_minus**2 - hbar * source_free_minus - c**2,
        "source-free minus Riccati residue",
    )
    require_zero(
        (source_free_plus - source_free_minus) / 2 - exact_residue,
        "exact source-free branch-difference residue",
    )
    exact_series = sp.series(exact_residue, hbar, 0, 6).removeO()
    generated_series = c + hbar**2 * residue_two + hbar**4 * residue_four
    require_zero(
        exact_series - generated_series,
        "double-pole residue versus exact square-root series",
    )

    # More generally, r(hbar)=r_0+hbar**2 r_2+... contributes r_2 and
    # the universal Riccati 1/4 in the same square root.
    r_zero = sp.symbols("r_0", positive=True)
    r_two = sp.symbols("r_2")
    generic_residue = sp.sqrt(
        r_zero + hbar**2 * (r_two + sp.Rational(1, 4))
    )
    generic_series = sp.series(generic_residue, hbar, 0, 4).removeO()
    expected_generic = (
        sp.sqrt(r_zero)
        + hbar**2
        * (r_two + sp.Rational(1, 4))
        / (2 * sp.sqrt(r_zero))
    )
    require_zero(
        generic_series - expected_generic,
        "generic order-hbar**2 double-pole residue",
    )

    # With r(hbar)=c**2-hbar**2/4, the two local Riccati residues are
    # hbar/2 +/- c.  Their branch difference is exactly c, so the regular
    # part Omega-lambda_0 has no logarithmic residue.
    langer_source = c**2 - hbar**2 / 4
    langer_plus = hbar / 2 + c
    langer_minus = hbar / 2 - c
    require_zero(
        langer_plus**2 - hbar * langer_plus - langer_source,
        "Langer plus Riccati residue",
    )
    require_zero(
        langer_minus**2 - hbar * langer_minus - langer_source,
        "Langer minus Riccati residue",
    )
    langer_even_residue = reduced((langer_plus - langer_minus) / 2)
    require_zero(
        langer_even_residue - c,
        "Langer cancellation in the branch-difference residue",
    )

    positive_loop_period = sp.pi * sp.I * (
        source_free_plus - source_free_minus
    )
    require_zero(
        positive_loop_period - 2 * sp.pi * sp.I * exact_residue,
        "double-pole positive-loop period",
    )

    return {
        "source_free_residue": exact_series,
        "lambda_two_residue": residue_two,
        "lambda_four_residue": residue_four,
        "langer_residue": langer_even_residue,
    }


def weber_quantum_period_audit() -> dict[str, object]:
    """Check closed and regularized relative Weber periods through hbar**6."""
    t = sp.symbols("t")
    a = sp.symbols("a", positive=True)
    z = a * (t + 1 / t) / 2
    y = a * (t - 1 / t) / 2
    dz_dt = sp.diff(z, t)

    require_zero(y**2 - (z**2 - a**2), "Weber rational parametrization")
    coefficients = generate_branch_on_cover(y, z, t, 6)
    one_forms = [reduced(coefficient * dz_dt) for coefficient in coefficients]

    expected_lambda_two = -(
        t * (3 * t**4 + 14 * t**2 + 3)
    ) / (2 * a**2 * (t - 1) ** 4 * (t + 1) ** 4)
    require_zero(
        one_forms[2] - expected_lambda_two,
        "printed Weber p_2 dz identity",
    )
    expected_primitive_two = (
        9 * t**4 + 12 * t**2 - 1
    ) / (12 * a**2 * (t - 1) ** 3 * (t + 1) ** 3)
    require_zero(
        sp.diff(expected_primitive_two, t) - one_forms[2],
        "printed Weber p_2 dz primitive",
    )

    classical_residue = local_residue_at_infinity(one_forms[0], t)
    opposite_classical_residue = reduced(sp.residue(one_forms[0], t, 0))
    require_zero(
        classical_residue - a**2 / 2,
        "Weber classical residue at infinity_+",
    )
    require_zero(
        opposite_classical_residue + a**2 / 2,
        "Weber classical residue at infinity_-",
    )
    require_zero(
        classical_residue + opposite_classical_residue,
        "sum of the Weber infinity residues",
    )
    closed_period = reduced(2 * sp.pi * sp.I * classical_residue)
    require_zero(
        closed_period - sp.I * sp.pi * a**2,
        "Weber oriented closed classical period",
    )

    relative_coefficients: list[sp.Expr] = []
    for k in range(1, 4):
        order = 2 * k
        one_form = one_forms[order]

        # Ramification poles carry no residue, and the quantum corrections
        # are residue-free at both points above infinity through this order.
        for point, point_label in ((1, "+1"), (-1, "-1"), (0, "0")):
            require_zero(
                sp.residue(one_form, t, point),
                f"Weber lambda_{order} residue at t={point_label}",
            )
        require_zero(
            local_residue_at_infinity(one_form, t),
            f"Weber lambda_{order} residue at t=infinity",
        )

        primitive = reduced(sp.integrate(one_form, t))
        require_zero(
            sp.diff(primitive, t) - one_form,
            f"Weber lambda_{order} rational primitive",
        )
        endpoint_value = reduced(
            sp.limit(primitive, t, sp.oo) - sp.limit(primitive, t, 0)
        )
        bernoulli_value = reduced(
            (1 - 2 ** (2 * k - 1))
            * sp.bernoulli(2 * k)
            / (2 * k * (2 * k - 1) * a ** (4 * k - 2))
        )
        require_zero(
            endpoint_value - bernoulli_value,
            f"Weber relative coefficient at hbar^{order}",
        )
        relative_coefficients.append(endpoint_value)

    expected_values = [
        -1 / (12 * a**2),
        sp.Rational(7, 360) / a**6,
        -sp.Rational(31, 1260) / a**10,
    ]
    for order, actual, expected in zip(
        (2, 4, 6), relative_coefficients, expected_values
    ):
        require_zero(
            actual - expected,
            f"explicit Weber endpoint coefficient hbar^{order}",
        )

    return {
        "closed_period": closed_period,
        "relative_coefficients": relative_coefficients,
        "lambda_two": one_forms[2],
    }


def schwarzian_covariance_audit() -> dict[str, sp.Expr]:
    """Check that the Schwarzian removes a spurious transformed p_2."""
    w = sp.symbols("w")
    kappa = sp.symbols("kappa", nonzero=True)
    z = sp.exp(w)
    slope = sp.diff(z, w)
    schwarzian = reduced(
        sp.diff(z, w, 3) / slope
        - sp.Rational(3, 2) * (sp.diff(z, w, 2) / slope) ** 2
    )
    require_zero(
        schwarzian + sp.Rational(1, 2),
        "Schwarzian of z=exp(w)",
    )

    transformed_leading = kappa * slope
    transformed_r_two = reduced(-schwarzian / 2)
    require_zero(
        transformed_r_two - sp.Rational(1, 4),
        "coordinate-induced order-hbar**2 source",
    )

    covariant_p_two = reduced(
        transformed_r_two / (2 * transformed_leading)
        + sp.diff(transformed_leading, w, 2)
        / (4 * transformed_leading**2)
        - 3
        * sp.diff(transformed_leading, w) ** 2
        / (8 * transformed_leading**3)
    )
    require_zero(
        covariant_p_two,
        "Schwarzian cancellation in transformed p_2",
    )

    naive_p_two = reduced(
        sp.diff(transformed_leading, w, 2)
        / (4 * transformed_leading**2)
        - 3
        * sp.diff(transformed_leading, w) ** 2
        / (8 * transformed_leading**3)
    )
    expected_naive = -sp.exp(-w) / (8 * kappa)
    require_zero(
        naive_p_two - expected_naive,
        "spurious p_2 when the Schwarzian is omitted",
    )

    return {
        "schwarzian": schwarzian,
        "transformed_r_two": transformed_r_two,
        "covariant_p_two": covariant_p_two,
        "naive_p_two": naive_p_two,
    }


def main() -> None:
    """Run every audit and print a compact reproducibility ledger."""
    print("Advanced Linear ODEs -- regularized quantum-period audit")
    print(f"Python: {platform.python_version()}")
    print(f"SymPy: {sp.__version__}")
    print()

    airy = airy_turning_point_audit()
    print("[PASS] Airy turning-point half-contour normalization")
    print(f"  lambda_2 = ({airy['lambda_two']}) dxi")
    print(f"  half-contour correction = {airy['half_quantum']}")

    double_pole = double_pole_audit()
    print("[PASS] Double-pole residues and Langer cancellation")
    print(f"  source-free residue = {double_pole['source_free_residue']} + O(hbar**6)")
    print(f"  Langer-corrected residue = {double_pole['langer_residue']}")

    weber = weber_quantum_period_audit()
    print("[PASS] Weber closed and regularized relative periods")
    print(f"  closed period = {weber['closed_period']}")
    print(f"  relative coefficients = {weber['relative_coefficients']}")

    covariance = schwarzian_covariance_audit()
    print("[PASS] Coordinate covariance under z=exp(w)")
    print(f"  Schwarzian = {covariance['schwarzian']}")
    print(f"  covariant p_2 = {covariance['covariant_p_two']}")
    print(f"  naive p_2 = {covariance['naive_p_two']}")

    print()
    print("All exact quantum-period checks passed.")
    print("Scope: formal coefficients only; no Borel or quantization claim tested.")


if __name__ == "__main__":
    main()