#!/usr/bin/env python3
"""Exact audit of the classical WKB action formulas on Chapter 8, Page 4.

The page uses

    hbar**2 psi'' = R psi,       lambda_0 = y dz,       y**2 = R_0.

This script checks only the algebraic and analytic calculations that can be
made explicit without pretending to prove homological statements by computer:

1. the Airy pullback, primitive, and sheet sign;
2. the Weber semicircle integral, one-way action, and doubled cut-cycle action;
3. the Laurent expansions and opposite residues at the two Weber infinities;
4. the Weber instance of the Gauss--Manin derivative formula; and
5. several integral 2 x 2 changes of symplectic basis.

Every test raises ``RuntimeError`` explicitly.  The audit is therefore unchanged
by Python's ``-O`` flag.  No Borel summation, Stokes phenomenon, endpoint
regularization, or global topological classification is asserted here.
"""

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(condition: bool, label: str, detail: object = "") -> None:
    """Raise a visible failure instead of relying on ``assert``."""
    if not condition:
        suffix = f": {detail}" if detail != "" else ""
        raise RuntimeError(f"{label} failed{suffix}")


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


def require_matrix_zero(matrix: sp.Matrix, label: str) -> None:
    """Raise unless every entry of a symbolic matrix vanishes."""
    residual = matrix.applyfunc(sp.simplify)
    if residual != sp.zeros(*residual.shape):
        raise RuntimeError(f"{label} failed: residual = {residual}")


def airy_audit() -> dict[str, sp.Expr]:
    """Check the Airy cover coordinate and its classical primitive."""
    xi = sp.symbols("xi")
    z = sp.symbols("z", positive=True)

    z_of_xi = xi**2
    dz_dxi = sp.diff(z_of_xi, xi)
    pulled_quadratic_coefficient = z_of_xi * dz_dxi**2
    lambda_xi_coefficient = 2 * xi**2
    primitive_xi = sp.Rational(2, 3) * xi**3

    require_zero(
        lambda_xi_coefficient**2 - pulled_quadratic_coefficient,
        "Airy lambda_0 squared equals the pulled-back quadratic differential",
    )
    require_zero(
        sp.diff(primitive_xi, xi) - lambda_xi_coefficient,
        "Airy primitive in the cover coordinate",
    )

    primitive_z = sp.Rational(2, 3) * z ** sp.Rational(3, 2)
    require_zero(
        sp.diff(primitive_z, z) - sp.sqrt(z),
        "Airy primitive on the chosen sheet",
    )
    require_zero(
        sp.diff(-primitive_z, z) + sp.sqrt(z),
        "Airy primitive changes sign on the other sheet",
    )

    return {
        "pullback_lambda": lambda_xi_coefficient,
        "primitive_xi": primitive_xi,
        "primitive_z": primitive_z,
    }


def weber_actions_audit() -> dict[str, sp.Expr]:
    """Evaluate the Weber semicircle, open action, and doubled action."""
    a = sp.symbols("a", positive=True)
    theta = sp.symbols("theta", real=True)

    z_theta = a * sp.sin(theta)
    y_theta = sp.I * a * sp.cos(theta)
    dz_dtheta = sp.diff(z_theta, theta)

    require_zero(
        y_theta**2 - (z_theta**2 - a**2),
        "Weber trigonometric parametrization lies on the spectral curve",
    )

    lower = -sp.pi / 2
    upper = sp.pi / 2
    semicircle_area = sp.integrate(
        a**2 * sp.cos(theta) ** 2,
        (theta, lower, upper),
    )
    require_zero(
        semicircle_area - sp.pi * a**2 / 2,
        "Weber semicircle area",
    )

    open_action = sp.integrate(
        y_theta * dz_dtheta,
        (theta, lower, upper),
    )
    expected_open = sp.I * sp.pi * a**2 / 2
    require_zero(open_action - expected_open, "Weber one-way action")

    other_sheet_open = sp.integrate(
        -y_theta * dz_dtheta,
        (theta, lower, upper),
    )
    require_zero(
        other_sheet_open + open_action,
        "Weber action changes sign under the deck involution",
    )

    # For delta_beta = beta - tau_* beta, the two lifted arcs have the same
    # endpoint orientation.  Anti-invariance therefore doubles the open action.
    doubled_action = reduced(open_action - other_sheet_open)
    expected_doubled = sp.I * sp.pi * a**2
    require_zero(
        doubled_action - expected_doubled,
        "Weber doubled cut-cycle action",
    )
    require_zero(
        doubled_action - 2 * open_action,
        "closed action equals twice the one-way action",
    )

    return {
        "a": a,
        "semicircle_area": semicircle_area,
        "open_action": open_action,
        "doubled_action": doubled_action,
    }


def coefficient(expression: sp.Expr, variable: sp.Symbol, power: int) -> sp.Expr:
    """Extract one Laurent coefficient by a residue formula."""
    return sp.residue(expression / variable ** (power + 1), variable, 0)


def weber_infinity_audit(a: sp.Symbol) -> dict[str, sp.Expr]:
    """Check both Laurent branches of lambda_0 at Weber infinity."""
    w = sp.symbols("w")

    # At infinity_+, y ~ z.  With z=1/w and dz=-w**(-2) dw,
    # the displayed expressions are the coefficients of dw.
    lambda_plus = -sp.sqrt(1 - a**2 * w**2) / w**3
    lambda_minus = -lambda_plus
    expected_plus = (
        -w**-3
        + a**2 * w**-1 / 2
        + a**4 * w / 8
        + a**6 * w**3 / 16
    )
    plus_series = sp.series(lambda_plus, w, 0, 5).removeO()
    minus_series = sp.series(lambda_minus, w, 0, 5).removeO()

    require_zero(
        plus_series - expected_plus,
        "Weber Laurent expansion at infinity_+ through w^3",
    )
    require_zero(
        minus_series + expected_plus,
        "Weber Laurent expansion at infinity_- through w^3",
    )
    require_zero(
        lambda_plus + lambda_minus,
        "Weber infinity branches are opposite",
    )

    residue_plus = coefficient(lambda_plus, w, -1)
    residue_minus = coefficient(lambda_minus, w, -1)
    require_zero(
        residue_plus - a**2 / 2,
        "Weber residue at infinity_+",
    )
    require_zero(
        residue_minus + a**2 / 2,
        "Weber residue at infinity_-",
    )
    require_zero(
        residue_plus + residue_minus,
        "sum of the two Weber residues",
    )

    positive_loop_period = 2 * sp.pi * sp.I * residue_plus
    require_zero(
        positive_loop_period - sp.I * sp.pi * a**2,
        "positive infinity_+ loop period",
    )

    return {
        "plus_series": plus_series,
        "minus_series": minus_series,
        "residue_plus": residue_plus,
        "residue_minus": residue_minus,
        "positive_loop_period": positive_loop_period,
    }


def gauss_manin_audit() -> dict[str, sp.Expr]:
    """Check d Pi/dE = -(1/2) integral dz/y for the Weber family."""
    E = sp.symbols("E", positive=True)
    z = sp.symbols("z")
    theta = sp.symbols("theta", real=True)
    a = sp.symbols("a", positive=True)

    y_formal = sp.sqrt(z**2 - E)
    require_zero(
        sp.diff(y_formal, E) + 1 / (2 * y_formal),
        "pointwise Gauss--Manin derivative of y dz",
    )

    # On beta, z=a sin(theta) and y=+i a cos(theta).  The doubled integral
    # of dz/y represents the same anti-invariant cut cycle as the action.
    z_theta = a * sp.sin(theta)
    y_theta = sp.I * a * sp.cos(theta)
    dz_over_y = sp.cancel(sp.diff(z_theta, theta) / y_theta)
    require_zero(dz_over_y + sp.I, "Weber dz/y pullback")

    beta_integral = sp.integrate(
        dz_over_y,
        (theta, -sp.pi / 2, sp.pi / 2),
    )
    closed_integral = 2 * beta_integral
    require_zero(
        beta_integral + sp.I * sp.pi,
        "Weber one-way integral of dz/y",
    )
    require_zero(
        closed_integral + 2 * sp.pi * sp.I,
        "Weber closed integral of dz/y",
    )

    period = sp.I * sp.pi * E
    period_derivative = sp.diff(period, E)
    gauss_manin_rhs = -closed_integral / 2
    require_zero(
        period_derivative - gauss_manin_rhs,
        "Weber Gauss--Manin period derivative",
    )

    open_period = sp.I * sp.pi * E / 2
    require_zero(
        sp.diff(open_period, E) + beta_integral / 2,
        "Weber open-action derivative",
    )

    return {
        "beta_integral": beta_integral,
        "closed_integral": closed_integral,
        "period_derivative": period_derivative,
        "gauss_manin_rhs": gauss_manin_rhs,
    }


def symplectic_basis_audit() -> dict[str, sp.Matrix]:
    """Check representative integral changes of a genus-one cycle basis."""
    intersection = sp.Matrix([[0, 1], [-1, 0]])
    matrices = {
        "identity": sp.eye(2),
        "positive twist about A": sp.Matrix([[1, 0], [1, 1]]),
        "positive twist about B": sp.Matrix([[1, -1], [0, 1]]),
        "quarter turn": sp.Matrix([[0, 1], [-1, 0]]),
        "integral composite": sp.Matrix([[2, 1], [1, 1]]),
    }
    period_a, period_b = sp.symbols("Pi_A Pi_B")
    period_column = sp.Matrix([period_a, period_b])

    for label, matrix in matrices.items():
        require(matrix.shape == (2, 2), f"{label}: matrix shape", matrix.shape)
        require(
            all(entry.is_integer is True for entry in matrix),
            f"{label}: integral entries",
            matrix,
        )
        require_zero(matrix.det() - 1, f"{label}: determinant one")
        require_matrix_zero(
            matrix * intersection * matrix.T - intersection,
            f"{label}: transformed cycle intersection form",
        )
        require_matrix_zero(
            matrix.T * intersection * matrix - intersection,
            f"{label}: symplectic matrix identity",
        )
        inverse = matrix.inv()
        require(
            all(entry.is_integer is True for entry in inverse),
            f"{label}: integral inverse",
            inverse,
        )

        transformed_periods = matrix * period_column
        for row in range(2):
            expected = sum(
                matrix[row, column] * period_column[column, 0]
                for column in range(2)
            )
            require_zero(
                transformed_periods[row, 0] - expected,
                f"{label}: period column follows the cycle basis",
            )

    orientation_reversing_swap = sp.Matrix([[0, 1], [1, 0]])
    require_zero(
        orientation_reversing_swap.det() + 1,
        "plain A/B swap has determinant minus one",
    )
    require_matrix_zero(
        orientation_reversing_swap
        * intersection
        * orientation_reversing_swap.T
        + intersection,
        "plain A/B swap reverses the intersection orientation",
    )

    return matrices


def main() -> None:
    """Run every exact audit and print a compact reproducibility ledger."""
    print("Advanced Linear ODEs -- WKB homology/action audit")
    print(f"Python: {platform.python_version()}")
    print(f"SymPy: {sp.__version__}")
    print()

    airy = airy_audit()
    print("[PASS] Airy cover and primitive")
    print(f"  lambda_0 = ({airy['pullback_lambda']}) dxi")
    print(f"  integral_0^z lambda_0 = {airy['primitive_z']}")

    weber = weber_actions_audit()
    print("[PASS] Weber semicircle and cut-cycle actions")
    print(f"  semicircle area = {weber['semicircle_area']}")
    print(f"  open action = {weber['open_action']}")
    print(f"  doubled action = {weber['doubled_action']}")

    infinity = weber_infinity_audit(weber["a"])
    print("[PASS] Weber expansions and opposite infinity residues")
    print(f"  infinity_+ series = {infinity['plus_series']} + O(w^5)")
    print(f"  residues = ({infinity['residue_plus']}, {infinity['residue_minus']})")
    print(f"  positive loop period = {infinity['positive_loop_period']}")

    gauss_manin = gauss_manin_audit()
    print("[PASS] Weber Gauss--Manin derivative integral")
    print(f"  closed integral dz/y = {gauss_manin['closed_integral']}")
    print(f"  dPi/dE = {gauss_manin['period_derivative']}")

    matrices = symplectic_basis_audit()
    print("[PASS] Integral 2 x 2 symplectic basis changes")
    for label, matrix in matrices.items():
        print(f"  {label}: {matrix.tolist()}")

    print()
    print("All algebraic and analytic checks passed.")
    print("Scope: no topological, regularization, or resummation claim was tested.")


if __name__ == "__main__":
    main()