#!/usr/bin/env python3
"""Exact audit of the formal WKB recursion and branch decomposition.

Dependencies used for the published calculation:
    Python 3.10.16
    SymPy 1.13.1

The convention is

    [-hbar**2 d_z**2 + V(z)] psi(z) = E psi(z),
    R(z) = V(z) - E,
    P = hbar psi'/psi,

so that P**2 + hbar P' = R.  On a local chart with a chosen square root
rho(z)**2 = R(z), this program:

1. generates both Riccati branches through hbar**6 independently;
2. verifies every Riccati coefficient through that order;
3. verifies P_-(z,hbar) = -P_+(z,-hbar);
4. checks that P_even=(P_+-P_-)/2 has only even powers and that
   P_amp=(P_++P_-)/2 has only odd powers;
5. compares P_3 with its printed total-derivative form and verifies

       P_amp = -(hbar/2) d_z log(P_even)

   coefficient by coefficient with explicit truncation bookkeeping;
6. verifies the reduced nonlinear equation for P_even coefficientwise;
7. contrasts an even hbar-dependent source, which preserves branch parity,
   with an odd source, which breaks it in a controlled way;
8. checks the normalized formal wavefunctions through hbar**6 and audits
   the sign and normalization of their ordered formal Wronskian; and
9. checks phase, amplitude, and Wronskian covariance for z=w**2; and
10. specializes the generic coefficients to the Airy potential R(z)=z,
   compares them with an independent exact coefficient table, and checks
   the exact Airy solution and the simple-turning constants A_1,...,A_6.

All checks raise explicit RuntimeError exceptions, so running Python with
-O does not disable the audit.  The calculation is formal and local: rho
must be nonzero on the chart used for the recursion.
"""

from __future__ import annotations

import platform
import sys

import sympy as sp


MAX_ORDER = 6


def reduced(expression: sp.Expr) -> sp.Expr:
    """Put a rational differential expression over a common denominator."""
    return sp.factor(sp.cancel(sp.together(expression)))


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


def generate_branch(
    rho: sp.Expr,
    sign: int,
    z: sp.Symbol,
    order: int,
) -> list[sp.Expr]:
    """Generate one branch of P**2+hbar*P'=rho**2 recursively."""
    if sign not in (-1, 1):
        raise RuntimeError(f"branch sign must be +1 or -1, got {sign}")
    if order < 0:
        raise RuntimeError(f"formal order must be nonnegative, got {order}")

    coefficients = [sign * rho]
    for n in range(1, order + 1):
        convolution = sum(
            coefficients[j] * coefficients[n - j]
            for j in range(1, n)
        )
        coefficient = -(
            sp.diff(coefficients[n - 1], z) + convolution
        ) / (2 * coefficients[0])
        coefficients.append(reduced(coefficient))
    return coefficients


def riccati_coefficient_residual(
    coefficients: list[sp.Expr],
    radicand: sp.Expr,
    z: sp.Symbol,
    n: int,
) -> sp.Expr:
    """Return [hbar**n](P**2+hbar*P'-R)."""
    if n < 0 or n >= len(coefficients):
        raise RuntimeError(
            f"coefficient index {n} lies outside 0..{len(coefficients) - 1}"
        )
    residual = sum(
        coefficients[j] * coefficients[n - j] for j in range(n + 1)
    )
    if n == 0:
        residual -= radicand
    else:
        residual += sp.diff(coefficients[n - 1], z)
    return reduced(residual)


def generate_sourced_branch(
    source_coefficients: list[sp.Expr],
    leading_root: sp.Expr,
    sign: int,
    z: sp.Symbol,
    order: int,
) -> list[sp.Expr]:
    """Generate a branch when R itself is a formal hbar series."""
    if sign not in (-1, 1):
        raise RuntimeError(f"branch sign must be +1 or -1, got {sign}")
    if order < 0:
        raise RuntimeError(f"formal order must be nonnegative, got {order}")
    if len(source_coefficients) <= order:
        raise RuntimeError(
            "the sourced recurrence needs R_0 through "
            f"R_{order}, got only {len(source_coefficients)} coefficients"
        )

    coefficients = [sign * leading_root]
    for n in range(1, order + 1):
        convolution = sum(
            coefficients[j] * coefficients[n - j]
            for j in range(1, n)
        )
        coefficient = (
            source_coefficients[n]
            - sp.diff(coefficients[n - 1], z)
            - convolution
        ) / (2 * coefficients[0])
        coefficients.append(reduced(coefficient))
    return coefficients


def sourced_riccati_coefficient_residual(
    coefficients: list[sp.Expr],
    source_coefficients: list[sp.Expr],
    z: sp.Symbol,
    n: int,
) -> sp.Expr:
    """Return [hbar**n](P**2+hbar*P'-R(hbar))."""
    if n < 0 or n >= len(coefficients):
        raise RuntimeError(
            f"coefficient index {n} lies outside 0..{len(coefficients) - 1}"
        )
    if n >= len(source_coefficients):
        raise RuntimeError(
            f"source coefficient R_{n} was not supplied to the audit"
        )
    residual = sum(
        coefficients[j] * coefficients[n - j] for j in range(n + 1)
    )
    if n > 0:
        residual += sp.diff(coefficients[n - 1], z)
    residual -= source_coefficients[n]
    return reduced(residual)


def generic_recursion_audit() -> dict[str, object]:
    """Generate and verify both generic local Riccati branches."""
    z = sp.symbols("z")
    radicand = sp.Function("R")(z)
    rho = sp.sqrt(radicand)
    plus = generate_branch(rho, 1, z, MAX_ORDER)
    minus = generate_branch(rho, -1, z, MAX_ORDER)

    for branch_name, coefficients in (("+", plus), ("-", minus)):
        for n in range(MAX_ORDER + 1):
            require_zero(
                riccati_coefficient_residual(
                    coefficients, radicand, z, n
                ),
                f"generic P_{branch_name} Riccati coefficient hbar^{n}",
            )

    # This comparison uses a minus branch generated independently from the
    # recurrence, rather than defining it by the claimed parity relation.
    for n in range(MAX_ORDER + 1):
        expected_minus = (-1) ** (n + 1) * plus[n]
        require_zero(
            minus[n] - expected_minus,
            f"branch parity at hbar^{n}",
        )

    return {
        "z": z,
        "R": radicand,
        "rho": rho,
        "plus": plus,
        "minus": minus,
    }


def printed_p_three_audit(data: dict[str, object]) -> sp.Expr:
    """Compare generated P_3+ with the page's total derivative."""
    z = data["z"]
    rho = data["rho"]
    plus = data["plus"]
    if not isinstance(z, sp.Symbol) or not isinstance(rho, sp.Expr):
        raise RuntimeError("internal error: invalid P_3 comparison data")
    if not isinstance(plus, list) or len(plus) <= 3:
        raise RuntimeError("internal error: generated P_3 is missing")

    printed_primitive = (
        -sp.diff(rho, z, 2) / (8 * rho**3)
        + 3 * sp.diff(rho, z) ** 2 / (16 * rho**4)
    )
    printed_p_three = sp.diff(printed_primitive, z)
    require_zero(
        plus[3] - printed_p_three,
        "generated P_3+ versus printed total derivative",
    )
    require_zero(
        printed_p_three + sp.diff(plus[2] / plus[0], z) / 2,
        "P_3+=-(1/2)d_z(P_2+/P_0+)",
    )
    return reduced(plus[3] - printed_p_three)


def formal_quotient_coefficients(
    numerator_coefficients: list[sp.Expr],
    denominator_coefficients: list[sp.Expr],
    order: int,
) -> list[sp.Expr]:
    """Return coefficients of one formal series divided by another."""
    if not denominator_coefficients:
        raise RuntimeError("the denominator series must have a leading term")
    available = min(
        len(numerator_coefficients), len(denominator_coefficients)
    )
    if order < 0 or order >= available:
        raise RuntimeError(
            "formal-quotient order must lie between zero and "
            f"{available - 1}, got {order}"
        )

    quotient: list[sp.Expr] = []
    leading = denominator_coefficients[0]
    for n in range(order + 1):
        lower_products = sum(
            denominator_coefficients[j] * quotient[n - j]
            for j in range(1, n + 1)
        )
        q_n = (numerator_coefficients[n] - lower_products) / leading
        quotient.append(reduced(q_n))
    return quotient


def logarithmic_derivative_coefficients(
    series_coefficients: list[sp.Expr],
    z: sp.Symbol,
    order: int,
) -> list[sp.Expr]:
    """Return coefficients of A'/A through the requested formal order."""
    numerator = [sp.diff(coefficient, z) for coefficient in series_coefficients]
    return formal_quotient_coefficients(
        numerator, series_coefficients, order
    )


def decomposition_audit(
    data: dict[str, object],
) -> dict[str, object]:
    """Check the even/odd powers and the exact amplitude identity."""
    z = data["z"]
    plus = data["plus"]
    minus = data["minus"]
    if not isinstance(z, sp.Symbol):
        raise RuntimeError("internal error: z is not a SymPy symbol")
    if not isinstance(plus, list) or not isinstance(minus, list):
        raise RuntimeError("internal error: branch coefficients are missing")

    p_even = [reduced((a - b) / 2) for a, b in zip(plus, minus)]
    p_amp = [reduced((a + b) / 2) for a, b in zip(plus, minus)]

    for n in range(MAX_ORDER + 1):
        if n % 2 == 1:
            require_zero(p_even[n], f"P_even odd-power coefficient hbar^{n}")
        else:
            require_zero(p_amp[n], f"P_amp even-power coefficient hbar^{n}")

    # To compare through hbar**MAX_ORDER, hbar*(P_even'/P_even) needs the
    # quotient only through hbar**(MAX_ORDER-1).  Its coefficients q_n are
    # obtained by formal division, not by asking SymPy for an implicit
    # infinite series.  Thus the checked hbar**n equation is
    #
    #     (P_amp)_n + q_{n-1}/2 = 0,       1 <= n <= MAX_ORDER.
    log_derivative = logarithmic_derivative_coefficients(
        p_even, z, MAX_ORDER - 1
    )
    require_zero(p_amp[0], "amplitude identity at hbar^0")
    for n in range(1, MAX_ORDER + 1):
        require_zero(
            p_amp[n] + log_derivative[n - 1] / 2,
            f"P_amp logarithmic identity at hbar^{n}",
        )

    return {
        "P_even": p_even,
        "P_amp": p_amp,
        "d_log_P_even": log_derivative,
        "identity order": MAX_ORDER,
        "quotient order required": MAX_ORDER - 1,
    }


def reduced_even_equation_audit(
    data: dict[str, object],
    decomposition: dict[str, object],
) -> int:
    """Verify the eliminated nonlinear equation for P_even."""
    z = data["z"]
    radicand = data["R"]
    p_even = decomposition["P_even"]
    log_derivative = decomposition["d_log_P_even"]
    if not isinstance(z, sp.Symbol) or not isinstance(radicand, sp.Expr):
        raise RuntimeError("internal error: invalid reduced-equation data")
    if not isinstance(p_even, list) or not isinstance(log_derivative, list):
        raise RuntimeError("internal error: P_even coefficient data missing")

    # The printed equation is
    #
    #   R = P_even^2
    #       -(hbar^2/2) P_even''/P_even
    #       +(3 hbar^2/4)(P_even'/P_even)^2.
    #
    # Through hbar**6, the two quotients are needed only through hbar**4.
    ratio_order = MAX_ORDER - 2
    second_derivative_numerator = [
        sp.diff(coefficient, z, 2) for coefficient in p_even
    ]
    second_derivative_ratio = formal_quotient_coefficients(
        second_derivative_numerator, p_even, ratio_order
    )

    for n in range(MAX_ORDER + 1):
        residual = sum(
            p_even[j] * p_even[n - j] for j in range(n + 1)
        )
        if n == 0:
            residual -= radicand
        if n >= 2:
            shifted = n - 2
            log_square = sum(
                log_derivative[j] * log_derivative[shifted - j]
                for j in range(shifted + 1)
            )
            residual += (
                -second_derivative_ratio[shifted] / 2
                + 3 * log_square / 4
            )
        require_zero(
            residual,
            f"reduced P_even equation coefficient hbar^{n}",
        )

    return MAX_ORDER


def sourced_parity_audit() -> dict[str, sp.Expr | int]:
    """Contrast even and odd hbar-dependent Riccati sources."""
    z = sp.symbols("z")
    rho = z + 2
    r_zero = rho**2
    zero = sp.Integer(0)

    # An even source obeys R(z,-hbar)=R(z,hbar), so the independently
    # generated branches must still satisfy P_-(hbar)=-P_+(-hbar).
    # Here R=(z+2)^2+hbar^2(z+1).
    even_source = [r_zero, zero, z + 1, zero, zero, zero, zero]
    even_plus = generate_sourced_branch(
        even_source, rho, 1, z, MAX_ORDER
    )
    even_minus = generate_sourced_branch(
        even_source, rho, -1, z, MAX_ORDER
    )
    require_zero(rho**2 - even_source[0], "even-source leading root")
    for branch_name, coefficients in (
        ("+", even_plus),
        ("-", even_minus),
    ):
        for n in range(MAX_ORDER + 1):
            require_zero(
                sourced_riccati_coefficient_residual(
                    coefficients, even_source, z, n
                ),
                f"even-source P_{branch_name} coefficient hbar^{n}",
            )
    even_part = [
        reduced((a - b) / 2) for a, b in zip(even_plus, even_minus)
    ]
    even_amplitude = [
        reduced((a + b) / 2) for a, b in zip(even_plus, even_minus)
    ]
    for n in range(MAX_ORDER + 1):
        require_zero(
            even_minus[n] - (-1) ** (n + 1) * even_plus[n],
            f"even-source branch parity at hbar^{n}",
        )
        if n % 2 == 1:
            require_zero(
                even_part[n],
                f"even-source P_even odd coefficient hbar^{n}",
            )
        else:
            require_zero(
                even_amplitude[n],
                f"even-source P_amp even coefficient hbar^{n}",
            )

    expected_even_p_two = (
        (z + 1) / (2 * (z + 2))
        - sp.Rational(3, 8) / (z + 2) ** 3
    )
    require_zero(
        even_plus[2] - expected_even_p_two,
        "even-source stated P_2+",
    )

    # An odd source no longer obeys hbar parity.  For
    # R=(z+2)^2+hbar(z+1), the hbar coefficient of the branch-antisymmetric
    # part is nonzero and has the displayed simple form.
    odd_source = [r_zero, z + 1, zero, zero, zero, zero, zero]
    odd_plus = generate_sourced_branch(
        odd_source, rho, 1, z, MAX_ORDER
    )
    odd_minus = generate_sourced_branch(
        odd_source, rho, -1, z, MAX_ORDER
    )
    require_zero(rho**2 - odd_source[0], "odd-source leading root")
    for branch_name, coefficients in (("+", odd_plus), ("-", odd_minus)):
        for n in range(MAX_ORDER + 1):
            require_zero(
                sourced_riccati_coefficient_residual(
                    coefficients, odd_source, z, n
                ),
                f"odd-source P_{branch_name} coefficient hbar^{n}",
            )

    odd_even_part_hbar = reduced((odd_plus[1] - odd_minus[1]) / 2)
    expected_odd_even_part_hbar = (z + 1) / (2 * (z + 2))
    require_zero(
        odd_even_part_hbar - expected_odd_even_part_hbar,
        "odd-source [hbar]P_even",
    )
    if reduced(odd_even_part_hbar) == 0:
        raise RuntimeError(
            "odd-source parity audit failed: [hbar]P_even vanished"
        )

    # At hbar^1 the even-source branch law would demand P_1-=P_1+.
    # Record its exact nonzero failure for the odd source.
    odd_branch_parity_failure = reduced(odd_minus[1] - odd_plus[1])
    expected_parity_failure = -(z + 1) / (z + 2)
    require_zero(
        odd_branch_parity_failure - expected_parity_failure,
        "odd-source branch-parity failure at hbar^1",
    )
    if reduced(odd_branch_parity_failure) == 0:
        raise RuntimeError(
            "odd-source audit failed: the branch-parity defect vanished"
        )

    return {
        "orders checked": MAX_ORDER,
        "even-source P_2+": expected_even_p_two,
        "odd-source [hbar]P_even": expected_odd_even_part_hbar,
        "odd-source hbar^1 branch defect": expected_parity_failure,
    }


def wavefunction_audit(
    data: dict[str, object],
    decomposition: dict[str, object],
) -> dict[str, object]:
    """Audit normalized formal wavefunctions and their ordered Wronskian."""
    z = data["z"]
    radicand = data["R"]
    plus = data["plus"]
    minus = data["minus"]
    p_even = decomposition["P_even"]
    log_derivative = decomposition["d_log_P_even"]
    if not isinstance(z, sp.Symbol) or not isinstance(radicand, sp.Expr):
        raise RuntimeError("internal error: invalid generic WKB data")
    if not all(
        isinstance(item, list)
        for item in (plus, minus, p_even, log_derivative)
    ):
        raise RuntimeError("internal error: invalid formal coefficient arrays")

    # The page's standard branch-symmetric formal wavefunctions are
    #
    #   widehat_psi_+ = P_even^(-1/2)
    #                   exp(+hbar^(-1) int P_even dz),
    #   widehat_psi_- = P_even^(-1/2)
    #                   exp(-hbar^(-1) int P_even dz).
    #
    # Their logarithmic momenta are +/-P_even-(hbar/2)d log P_even.
    # The quotient is needed only through hbar**5 to determine these
    # momenta through hbar**6.
    logarithmic_plus = [p_even[0]]
    logarithmic_minus = [-p_even[0]]
    for n in range(1, MAX_ORDER + 1):
        amplitude_term = -log_derivative[n - 1] / 2
        logarithmic_plus.append(reduced(p_even[n] + amplitude_term))
        logarithmic_minus.append(reduced(-p_even[n] + amplitude_term))

    for n in range(MAX_ORDER + 1):
        require_zero(
            logarithmic_plus[n] - plus[n],
            f"widehat_psi_+ logarithmic momentum at hbar^{n}",
        )
        require_zero(
            logarithmic_minus[n] - minus[n],
            f"widehat_psi_- logarithmic momentum at hbar^{n}",
        )
        require_zero(
            riccati_coefficient_residual(
                logarithmic_plus, radicand, z, n
            ),
            f"widehat_psi_+ ODE coefficient hbar^{n}",
        )
        require_zero(
            riccati_coefficient_residual(
                logarithmic_minus, radicand, z, n
            ),
            f"widehat_psi_- ODE coefficient hbar^{n}",
        )

    # Define W[f,g]=f*g'-f'*g, with the + branch first.  The exponential
    # factors cancel in widehat_psi_+ widehat_psi_-=P_even^(-1), while
    # hbar(widehat_psi_-'/widehat_psi_- -
    #      widehat_psi_+'/widehat_psi_+)=-2 P_even.  Consequently the
    # page's standard symmetric prefactor gives -2/hbar.  For comparison,
    # define chi_+/-=widehat_psi_+/-/sqrt(2); that separate rescaled pair
    # has ordered Wronskian -1/hbar.
    hbar = sp.symbols("hbar", nonzero=True)
    p_even_truncated = sum(
        hbar**n * coefficient for n, coefficient in enumerate(p_even)
    )
    exact_log_ratio = sp.diff(p_even_truncated, z) / p_even_truncated
    exact_logarithmic_plus = (
        p_even_truncated - hbar * exact_log_ratio / 2
    )
    exact_logarithmic_minus = (
        -p_even_truncated - hbar * exact_log_ratio / 2
    )
    standard_wronskian = reduced(
        (exact_logarithmic_minus - exact_logarithmic_plus)
        / (hbar * p_even_truncated)
    )
    chi_wronskian = reduced(standard_wronskian / 2)
    require_zero(
        standard_wronskian + 2 / hbar,
        "ordered W[widehat_psi_+,widehat_psi_-]",
    )
    require_zero(
        chi_wronskian + 1 / hbar,
        "ordered W[chi_+,chi_-] after a separate 1/sqrt(2) rescaling",
    )

    return {
        "ODE order checked": MAX_ORDER,
        "Wronskian convention": "W[f,g] = f*g' - f'*g",
        "W[widehat_psi_+,widehat_psi_-]": standard_wronskian,
        "W[chi_+,chi_-]": chi_wronskian,
    }


def coordinate_covariance_audit() -> dict[str, sp.Expr | str]:
    """Check phase, amplitude, Riccati, and Wronskian laws for z=w^2."""
    w = sp.symbols("w", positive=True)
    hbar = sp.symbols("hbar", nonzero=True)
    z_of_w = w**2
    jacobian = sp.diff(z_of_w, w)
    jacobian_prime = sp.diff(jacobian, w)
    schwarzian = reduced(
        sp.diff(z_of_w, w, 3) / jacobian
        - sp.Rational(3, 2) * (jacobian_prime / jacobian) ** 2
    )
    require_zero(
        schwarzian + sp.Rational(3, 2) / w**2,
        "Schwarzian of z=w^2",
    )

    r_value = sp.symbols("R_composed")
    transformed_r = reduced(
        jacobian**2 * r_value - hbar**2 * schwarzian / 2
    )
    require_zero(
        transformed_r
        - (4 * w**2 * r_value + 3 * hbar**2 / (4 * w**2)),
        "transformed R for z=w^2",
    )

    # Verify the affine branch law against the transformed Riccati equation
    # for two independent abstract branches.  Along z(w), d_w P=s P_z.
    branch_values = sp.symbols("P_plus P_minus")
    branch_derivatives = sp.symbols("P_plus_z P_minus_z")
    correction = hbar * jacobian_prime / (2 * jacobian)
    transformed_branches: list[sp.Expr] = []
    for name, p_value, p_z in zip(
        ("+", "-"), branch_values, branch_derivatives
    ):
        transformed_p = jacobian * p_value - correction
        transformed_p_derivative = (
            jacobian_prime * p_value
            + jacobian**2 * p_z
            - sp.diff(correction, w)
        )
        residual = (
            transformed_p**2
            + hbar * transformed_p_derivative
            - transformed_r
        ).subs(p_z, (r_value - p_value**2) / hbar)
        require_zero(
            residual,
            f"transformed abstract Riccati branch {name}",
        )
        transformed_branches.append(transformed_p)

    p_plus, p_minus = branch_values
    transformed_plus, transformed_minus = transformed_branches
    p_even = (p_plus - p_minus) / 2
    p_amp = (p_plus + p_minus) / 2
    transformed_even = reduced(
        (transformed_plus - transformed_minus) / 2
    )
    transformed_amp = reduced(
        (transformed_plus + transformed_minus) / 2
    )
    require_zero(
        transformed_even - jacobian * p_even,
        "phase momentum coordinate law",
    )
    require_zero(
        transformed_amp - (jacobian * p_amp - correction),
        "amplitude connection coordinate law",
    )

    # The first equation is precisely invariance of P_even dz.  The second
    # is P_amp dz shifted by -(hbar/2)d log(s).  Their coefficient laws in
    # the w chart were just checked without assuming relations between the
    # two branch values.
    require_zero(
        transformed_even - p_even * sp.diff(z_of_w, w),
        "phase one-form invariance",
    )
    require_zero(
        transformed_amp
        - (
            p_amp * sp.diff(z_of_w, w)
            - hbar * sp.diff(sp.log(jacobian), w) / 2
        ),
        "amplitude one-form connection shift",
    )

    # Finally check W_w[s^(-1/2)f(z(w)),s^(-1/2)g(z(w))]
    # =W_z[f,g] for arbitrary functions.  Positivity of w fixes the local
    # square-root branch of s=2w and lets SymPy differentiate it directly.
    z = sp.symbols("z")
    f = sp.Function("f")
    g = sp.Function("g")
    transformed_f = jacobian ** (-sp.Rational(1, 2)) * f(z_of_w)
    transformed_g = jacobian ** (-sp.Rational(1, 2)) * g(z_of_w)
    wronskian_w = reduced(
        transformed_f * sp.diff(transformed_g, w)
        - sp.diff(transformed_f, w) * transformed_g
    )
    wronskian_z = f(z) * sp.diff(g(z), z) - sp.diff(f(z), z) * g(z)
    wronskian_z_composed = wronskian_z.subs(z, z_of_w)
    require_zero(
        wronskian_w - wronskian_z_composed,
        "inverse-half-density Wronskian invariance",
    )

    return {
        "map": "z=w^2 on w>0",
        "Schwarzian": schwarzian,
        "transformed R": transformed_r,
        "phase-law residual": reduced(
            transformed_even - jacobian * p_even
        ),
        "amplitude-law residual": reduced(
            transformed_amp - (jacobian * p_amp - correction)
        ),
        "Wronskian residual": reduced(
            wronskian_w - wronskian_z_composed
        ),
    }


def simple_turning_audit(data: dict[str, object]) -> dict[str, object]:
    """Generate A_n and match the simple-turning Airy coefficients."""
    z = data["z"]
    radicand = data["R"]
    plus = data["plus"]
    if not isinstance(z, sp.Symbol) or not isinstance(radicand, sp.Expr):
        raise RuntimeError("internal error: invalid turning-point data")
    if not isinstance(plus, list):
        raise RuntimeError("internal error: plus-branch coefficients missing")

    # Store A_0 only as an unused placeholder so list indices equal the
    # mathematical subscripts.  The stated recurrence begins at A_1.
    constants = [sp.Integer(0), -sp.Rational(1, 4)]
    for n in range(2, MAX_ORDER + 1):
        convolution = sum(
            constants[j] * constants[n - j] for j in range(1, n)
        )
        a_n = (
            sp.Rational(3 * n - 4, 4) * constants[n - 1]
            - convolution / 2
        )
        constants.append(reduced(a_n))

    for n in range(1, MAX_ORDER + 1):
        a_n = constants[n]
        if a_n == 0 or a_n.is_negative is not True:
            raise RuntimeError(
                f"simple-turning negativity failed at A_{n}={a_n}"
            )

        # For the Airy model R=z, one has a=c=1 and t=z, hence
        # p_n=A_n z^(-(3n-1)/2).  Compare against the generic R-recursion,
        # not against a second copy of the constant recurrence.
        generated_airy = reduced(plus[n].subs(radicand, z).doit())
        expected_airy = a_n * z ** (-sp.Rational(3 * n - 1, 2))
        require_zero(
            generated_airy - expected_airy,
            f"simple-turning A_{n} versus generated Airy P_{n}",
        )

    return {
        "orders checked": MAX_ORDER,
        "A_1 through A_6": constants[1 : MAX_ORDER + 1],
        "all nonzero and negative": True,
    }


def airy_audit(data: dict[str, object]) -> dict[str, object]:
    """Regress the generic recursion against R=z and the exact Airy Ai."""
    z = data["z"]
    radicand = data["R"]
    plus = data["plus"]
    if not isinstance(z, sp.Symbol) or not isinstance(radicand, sp.Expr):
        raise RuntimeError("internal error: invalid Airy substitution data")
    if not isinstance(plus, list):
        raise RuntimeError("internal error: plus-branch coefficients missing")

    # For R=z on the positive real chart, P_n=c_n z^((1-3n)/2).
    # These rational constants are recorded independently of the generic
    # differential-expression generator above.
    expected_constants = (
        sp.Integer(1),
        -sp.Rational(1, 4),
        -sp.Rational(5, 32),
        -sp.Rational(15, 64),
        -sp.Rational(1105, 2048),
        -sp.Rational(1695, 1024),
        -sp.Rational(414125, 65536),
    )
    airy_coefficients: list[sp.Expr] = []
    for n, expected_constant in enumerate(expected_constants):
        specialized = plus[n].subs(radicand, z).doit()
        specialized = reduced(specialized)
        expected = expected_constant * z ** sp.Rational(1 - 3 * n, 2)
        require_zero(
            specialized - expected,
            f"Airy formal coefficient P_{n}",
        )
        airy_coefficients.append(specialized)

    # SymPy knows Ai''(x)=x Ai(x), so this is also an exact, nonformal
    # regression of hbar**2 psi''=z psi with the correct hbar scaling.
    z_positive = sp.symbols("z_positive", positive=True)
    hbar = sp.symbols("hbar", positive=True)
    airy_solution = sp.airyai(z_positive / hbar ** sp.Rational(2, 3))
    airy_residual = sp.simplify(
        hbar**2 * sp.diff(airy_solution, z_positive, 2)
        - z_positive * airy_solution
    )
    require_zero(airy_residual, "exact Airy Ai solution")

    return {
        "R(z)": z,
        "P_0 through P_6": airy_coefficients,
        "exact Ai residual": airy_residual,
    }


def print_coefficients(title: str, coefficients: list[sp.Expr]) -> None:
    """Print a deterministic formal coefficient list."""
    print(title)
    for n, coefficient in enumerate(coefficients):
        print(f"  hbar^{n}: {coefficient}")


def main() -> None:
    """Run every exact audit and print the verified identities."""
    if sys.version_info < (3, 10):
        raise RuntimeError("Python 3.10 or newer is required")

    print("runtime")
    print(f"  Python = {platform.python_version()}")
    print(f"  SymPy = {sp.__version__}")
    print()

    generic = generic_recursion_audit()
    plus = generic["plus"]
    minus = generic["minus"]
    if not isinstance(plus, list) or not isinstance(minus, list):
        raise RuntimeError("internal error: generic branch construction failed")
    print("generic local recursion (R nonzero, chosen rho=sqrt(R))")
    print(f"  orders checked on each branch = 0..{MAX_ORDER}")
    print("  branch identity = P_-(z,hbar) = -P_+(z,-hbar)")
    print(f"  P_0+ = {plus[0]}")
    print(f"  P_1+ = {plus[1]}")
    print(f"  P_2+ = {plus[2]}")
    print(f"  P_0- = {minus[0]}")
    p_three_residual = printed_p_three_audit(generic)
    print(
        "  printed P_3+ total-derivative comparison residual = "
        f"{p_three_residual}"
    )
    print()

    decomposition = decomposition_audit(generic)
    print("branch decomposition")
    print("  P_even = (P_+ - P_-)/2: even hbar powers only")
    print("  P_amp  = (P_+ + P_-)/2: odd hbar powers only")
    print(
        "  P_amp = -(hbar/2) d_z log(P_even): "
        f"checked through hbar^{decomposition['identity order']}"
    )
    print(
        "  quotient d_z log(P_even) required only through "
        f"hbar^{decomposition['quotient order required']}"
    )
    reduced_even_order = reduced_even_equation_audit(
        generic, decomposition
    )
    print(
        "  reduced P_even equation checked coefficientwise through "
        f"hbar^{reduced_even_order}"
    )
    print()

    sourced = sourced_parity_audit()
    print("hbar-dependent source parity")
    print(
        "  even source R=(z+2)^2+hbar^2(z+1): branch and power "
        f"parity checked through hbar^{sourced['orders checked']}"
    )
    print(f"  even-source P_2+ = {sourced['even-source P_2+']}")
    print(
        "  odd source R=(z+2)^2+hbar(z+1): "
        "[hbar]P_even = "
        f"{sourced['odd-source [hbar]P_even']}"
    )
    print(
        "  odd-source branch-parity defect P_1- - P_1+ = "
        f"{sourced['odd-source hbar^1 branch defect']}"
    )
    print()

    wavefunctions = wavefunction_audit(generic, decomposition)
    print("normalized formal wavefunctions")
    print(
        "  ODE residual coefficients checked on each branch = "
        f"hbar^0..hbar^{wavefunctions['ODE order checked']}"
    )
    print(f"  ordered convention: {wavefunctions['Wronskian convention']}")
    print(
        "  widehat_psi_+/- use P_even^(-1/2): "
        "W[widehat_psi_+,widehat_psi_-] = "
        f"{wavefunctions['W[widehat_psi_+,widehat_psi_-]']}"
    )
    print(
        "  separate chi_+/-=widehat_psi_+/-/sqrt(2): "
        f"W[chi_+,chi_-] = {wavefunctions['W[chi_+,chi_-]']}"
    )
    print()

    covariance = coordinate_covariance_audit()
    print("coordinate covariance")
    print(f"  map = {covariance['map']}")
    print(f"  Schwarzian = {covariance['Schwarzian']}")
    print(f"  transformed R = {covariance['transformed R']}")
    print(f"  phase-law residual = {covariance['phase-law residual']}")
    print(
        "  amplitude-law residual = "
        f"{covariance['amplitude-law residual']}"
    )
    print(f"  Wronskian residual = {covariance['Wronskian residual']}")
    print()

    turning = simple_turning_audit(generic)
    print("simple-turning constants")
    turning_constants = turning["A_1 through A_6"]
    if not isinstance(turning_constants, list):
        raise RuntimeError("internal error: turning constants are missing")
    for n, constant in enumerate(turning_constants, start=1):
        print(f"  A_{n} = {constant}")
    print(
        "  matched generated Airy coefficients; all constants through "
        f"A_{turning['orders checked']} are nonzero and negative"
    )
    print()

    airy = airy_audit(generic)
    airy_coefficients = airy["P_0 through P_6"]
    if not isinstance(airy_coefficients, list):
        raise RuntimeError("internal error: Airy coefficient list missing")
    print("Airy regression (R=z, positive-real square-root chart)")
    print_coefficients("  plus-branch coefficients", airy_coefficients)
    print(f"  exact Ai ODE residual = {airy['exact Ai residual']}")
    print("\nall exact checks passed")


if __name__ == "__main__":
    main()