#!/usr/bin/env python3
"""Exact audit for the hbar-dependent normal form and Riccati equation.

Dependencies used for the published calculation:
    Python 3.10.16
    SymPy 1.13.1

Only exact SymPy algebra is used.  The script verifies:

1. Liouville removal of the first derivative from u'' + p u' + q u = 0;
2. reduction of hbar**2 psi'' = R psi to P**2 + hbar P' = R;
3. finite-order regression checks for the unsourced and sourced recurrences;
4. P_0, P_1, P_2 and the displayed R_1-dependent P_1 formula;
5. exact constant-R exponential and cosh solutions; and
6. inverse-half-density covariance, including the Schwarzian term, for
   z = exp(w).

All checks raise explicit RuntimeError exceptions, so running Python with -O
does not disable the audit.
"""

from __future__ import annotations

import platform
import sys

import sympy as sp


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


def liouville_audit() -> dict[str, sp.Expr]:
    """Verify the Liouville gauge for arbitrary coefficient functions."""
    z = sp.symbols("z")
    p = sp.Function("p")(z)
    q = sp.Function("q")(z)
    primitive = sp.Function("I_p")(z)
    psi = sp.Function("psi")(z)

    # I_p' = p and u = exp(-I_p/2) psi.
    gauge = sp.exp(-primitive / 2)
    u = gauge * psi
    original_left = sp.diff(u, z, 2) + p * sp.diff(u, z) + q * u
    reduced_left = sp.simplify(
        (original_left / gauge).subs(
            {
                sp.diff(primitive, z): p,
                sp.diff(primitive, z, 2): sp.diff(p, z),
            },
            simultaneous=True,
        )
    )

    potential = sp.simplify(q - sp.diff(p, z) / 2 - p**2 / 4)
    expected_left = sp.diff(psi, z, 2) + potential * psi
    require_zero(reduced_left - expected_left, "Liouville reduction")

    return {
        "gauge factor": gauge,
        "normal-form coefficient": potential,
        "reduced residual": sp.simplify(reduced_left - expected_left),
    }


def riccati_audit() -> dict[str, sp.Expr]:
    """Verify the logarithmic ansatz, recursion, and first WKB terms."""
    z, hbar = sp.symbols("z hbar", nonzero=True)
    action = sp.Function("A")(z)
    riccati_momentum = sp.Function("P")(z)
    wavefunction = sp.exp(action / hbar)

    logarithmic_second_derivative = sp.simplify(
        hbar**2 * sp.diff(wavefunction, z, 2) / wavefunction
    )
    logarithmic_second_derivative = logarithmic_second_derivative.subs(
        {
            sp.diff(action, z): riccati_momentum,
            sp.diff(action, z, 2): sp.diff(riccati_momentum, z),
        },
        simultaneous=True,
    )
    expected_riccati_left = riccati_momentum**2 + hbar * sp.diff(
        riccati_momentum, z
    )
    require_zero(
        logarithmic_second_derivative - expected_riccati_left,
        "exponential-ansatz reduction",
    )

    # Check the coefficient formula through a generic finite order.  Since
    # the check uses independent functions P_n(z), it tests the convolution
    # and the derivative-index shift without assuming the low-order answer.
    order = 5
    coefficients = [sp.Function(f"P_{n}")(z) for n in range(order + 1)]
    series = sum(
        hbar**n * coefficient for n, coefficient in enumerate(coefficients)
    )
    r_static = sp.Function("R")(z)
    riccati_residual = sp.expand(
        series**2 + hbar * sp.diff(series, z) - r_static
    )

    zeroth_expected = coefficients[0] ** 2 - r_static
    require_zero(
        riccati_residual.coeff(hbar, 0) - zeroth_expected,
        "Riccati recursion at order zero",
    )
    for n in range(1, order + 1):
        convolution = sum(
            coefficients[k] * coefficients[n - k] for k in range(1, n)
        )
        expected = (
            2 * coefficients[0] * coefficients[n]
            + sp.diff(coefficients[n - 1], z)
            + convolution
        )
        require_zero(
            riccati_residual.coeff(hbar, n) - expected,
            f"Riccati recursion at order {n}",
        )

    r_function = sp.Function("R")(z)
    r_prime = sp.diff(r_function, z)
    r_second = sp.diff(r_function, z, 2)
    p_zero_base = sp.sqrt(r_function)
    p_one = -r_prime / (4 * r_function)
    p_two_base = (
        4 * r_function * r_second - 5 * r_prime**2
    ) / (32 * r_function ** sp.Rational(5, 2))

    for sign, branch_name in ((1, "+"), (-1, "-")):
        p_zero = sign * p_zero_base
        p_two = sign * p_two_base
        require_zero(p_zero**2 - r_function, f"P_0{branch_name}")
        require_zero(
            2 * p_zero * p_one + sp.diff(p_zero, z),
            f"P_1 on the {branch_name} branch",
        )
        require_zero(
            2 * p_zero * p_two + sp.diff(p_one, z) + p_one**2,
            f"P_2{branch_name}",
        )

    return {
        "Riccati equation": expected_riccati_left,
        "P_0+": p_zero_base,
        "P_0-": -p_zero_base,
        "P_1": p_one,
        "P_2+": p_two_base,
        "P_2-": -p_two_base,
        "finite unsourced orders checked": sp.Integer(order),
    }


def sourced_recurrence_audit() -> dict[str, sp.Expr]:
    """Check the recurrence when R itself has an hbar expansion."""
    z, hbar = sp.symbols("z hbar", nonzero=True)
    order = 5
    momenta = [sp.Function(f"P_{n}")(z) for n in range(order + 1)]
    sources = [sp.Function(f"R_{n}")(z) for n in range(order + 1)]
    momentum_series = sum(
        hbar**n * momentum for n, momentum in enumerate(momenta)
    )
    source_series = sum(
        hbar**n * source for n, source in enumerate(sources)
    )
    residual = sp.expand(
        momentum_series**2
        + hbar * sp.diff(momentum_series, z)
        - source_series
    )

    require_zero(
        residual.coeff(hbar, 0) - (momenta[0] ** 2 - sources[0]),
        "sourced Riccati recurrence at order zero",
    )
    for n in range(1, order + 1):
        convolution = sum(
            momenta[j] * momenta[n - j] for j in range(1, n)
        )
        expected = (
            2 * momenta[0] * momenta[n]
            + sp.diff(momenta[n - 1], z)
            + convolution
            - sources[n]
        )
        require_zero(
            residual.coeff(hbar, n) - expected,
            f"sourced Riccati recurrence at order {n}",
        )

    # If R_1 is nonzero, the coefficient equation is
    # 2 P_0 P_1 + P_0' = R_1.  Check the displayed solution on both
    # square-root branches rather than relying on a formal solve call.
    r_zero = sp.Function("R_0")(z)
    r_one = sp.Function("R_1")(z)
    p_zero_base = sp.sqrt(r_zero)
    common_part = -sp.diff(r_zero, z) / (4 * r_zero)
    source_part = r_one / (2 * p_zero_base)
    displayed_p_one: dict[str, sp.Expr] = {}
    for sign, branch_name in ((1, "+"), (-1, "-")):
        p_zero = sign * p_zero_base
        p_one = common_part + sign * source_part
        require_zero(
            2 * p_zero * p_one + sp.diff(p_zero, z) - r_one,
            f"R_1-dependent P_1{branch_name}",
        )
        displayed_p_one[f"P_1{branch_name} with R_1"] = p_one

    result = {
        "order-zero coefficient audit residual": sp.simplify(
            momenta[0] ** 2 - sources[0]
            - residual.coeff(hbar, 0)
        ),
        "finite sourced orders checked": sp.Integer(order),
    }
    result.update(displayed_p_one)
    return result


def constant_r_audit() -> dict[str, sp.Expr]:
    """Check exact exponential and cosh solutions when R is constant."""
    z = sp.symbols("z")
    hbar, kappa = sp.symbols("hbar kappa", nonzero=True)
    residuals: dict[str, sp.Expr] = {}

    for sign, branch_name in ((1, "+"), (-1, "-")):
        wavefunction = sp.exp(sign * kappa * z / hbar)
        residual = sp.simplify(
            hbar**2 * sp.diff(wavefunction, z, 2) - kappa**2 * wavefunction
        )
        require_zero(residual, f"constant-R solution {branch_name}")

        logarithmic_derivative = sp.simplify(
            hbar * sp.diff(wavefunction, z) / wavefunction
        )
        require_zero(
            logarithmic_derivative - sign * kappa,
            f"constant-R Riccati branch {branch_name}",
        )
        residuals[f"solution residual {branch_name}"] = residual

    cosh_solution = sp.cosh(kappa * z / hbar)
    cosh_ode_residual = sp.simplify(
        hbar**2 * sp.diff(cosh_solution, z, 2)
        - kappa**2 * cosh_solution
    )
    require_zero(cosh_ode_residual, "constant-R cosh solution")

    p_psi = sp.simplify(
        hbar * sp.diff(cosh_solution, z) / cosh_solution
    )
    expected_p_psi = kappa * sp.tanh(kappa * z / hbar)
    require_zero(p_psi - expected_p_psi, "cosh logarithmic momentum")
    cosh_riccati_residual = sp.trigsimp(
        p_psi**2 + hbar * sp.diff(p_psi, z) - kappa**2
    )
    require_zero(cosh_riccati_residual, "cosh Riccati equation")

    residuals["cosh ODE residual"] = cosh_ode_residual
    residuals["cosh P_psi"] = p_psi
    residuals["cosh Riccati residual"] = cosh_riccati_residual
    return residuals


def schwarzian(function: sp.Expr, coordinate: sp.Symbol) -> sp.Expr:
    """Return the Schwarzian derivative {function, coordinate}."""
    first = sp.diff(function, coordinate)
    return sp.simplify(
        sp.diff(function, coordinate, 3) / first
        - sp.Rational(3, 2)
        * (sp.diff(function, coordinate, 2) / first) ** 2
    )


def coordinate_covariance_audit() -> dict[str, sp.Expr]:
    """Verify the inverse-half-density law for z = exp(w)."""
    # For real w, z'=exp(w) is positive and its square root is unambiguous.
    # The map is intentionally nonlinear and has a nonzero Schwarzian, so
    # this test is sensitive to both the sign and the factor 1/2.
    w = sp.symbols("w", real=True)
    hbar, kappa = sp.symbols("hbar kappa", positive=True)
    z_of_w = sp.exp(w)
    jacobian = sp.diff(z_of_w, w)
    map_schwarzian = schwarzian(z_of_w, w)

    # First verify the operator identity for arbitrary Psi and R:
    #
    #   (hbar^2 d_w^2 - R_tilde) phi
    #       = (z')^(3/2) (hbar^2 d_z^2 - R) Psi,
    #
    # where phi=(z')^(-1/2)Psi(z(w)).
    z = sp.symbols("z")
    psi_function = sp.Function("Psi")
    r_function = sp.Function("R")
    generic_phi = (
        jacobian ** (-sp.Rational(1, 2)) * psi_function(z_of_w)
    )
    generic_transformed_r = sp.simplify(
        jacobian**2 * r_function(z_of_w)
        - hbar**2 * map_schwarzian / 2
    )
    generic_transformed_residual = sp.expand(
        hbar**2 * sp.diff(generic_phi, w, 2)
        - generic_transformed_r * generic_phi
    )
    original_residual_composed = (
        hbar**2
        * sp.Subs(sp.diff(psi_function(z), z, 2), z, z_of_w)
        - r_function(z_of_w) * psi_function(z_of_w)
    )
    require_zero(
        generic_transformed_residual
        - jacobian ** sp.Rational(3, 2) * original_residual_composed,
        "inverse-half-density operator covariance",
    )

    # Then use R=kappa^2 and both exact exponentials as a direct solution
    # check.  This also produces a compact transformed potential to print.
    transformed_r = sp.simplify(
        jacobian**2 * kappa**2 - hbar**2 * map_schwarzian / 2
    )

    require_zero(
        map_schwarzian + sp.Rational(1, 2),
        "Schwarzian of z = exp(w)",
    )
    require_zero(
        transformed_r - (kappa**2 * sp.exp(2 * w) + hbar**2 / 4),
        "transformed potential",
    )
    branch_residuals: dict[str, sp.Expr] = {}
    for sign, branch_name in ((1, "+"), (-1, "-")):
        psi_composed = sp.exp(sign * kappa * z_of_w / hbar)
        phi = jacobian ** (-sp.Rational(1, 2)) * psi_composed
        transformed_residual = sp.simplify(
            hbar**2 * sp.diff(phi, w, 2) - transformed_r * phi
        )
        require_zero(
            transformed_residual,
            f"coordinate-covariance solution {branch_name}",
        )
        branch_residuals[f"solution residual {branch_name}"] = (
            transformed_residual
        )

    result = {
        "z(w)": z_of_w,
        "Schwarzian {z,w}": map_schwarzian,
        "transformed R": transformed_r,
        "operator covariance residual": sp.simplify(
            generic_transformed_residual
            - jacobian ** sp.Rational(3, 2) * original_residual_composed
        ),
    }
    result.update(branch_residuals)
    return result


def print_section(title: str, values: dict[str, sp.Expr]) -> None:
    """Print one compact, deterministic audit section."""
    print(title)
    for name, value in values.items():
        print(f"  {name} = {value}")


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()

    print_section("Liouville normal form", liouville_audit())
    print()
    print_section("Riccati expansion", riccati_audit())
    print()
    print_section("hbar-dependent sourced recurrence", sourced_recurrence_audit())
    print()
    print_section("constant-R benchmarks", constant_r_audit())
    print()
    print_section(
        "inverse-half-density coordinate covariance",
        coordinate_covariance_audit(),
    )
    print("\nall exact checks passed")


if __name__ == "__main__":
    main()