#!/usr/bin/env python3
"""Reproduce the geometry checks for Chapter 9, Page 3.

The convention is

    [-hbar**2 d_z**2 + V(z)] psi(z) = E psi(z),
    Q_0(z) = V(z) - E,
    phi_0 = Q_0(z) dz**2,
    y**2 = Q_0.

For a Borel or graph direction theta, the phased quadratic differential
is phi_theta = exp(-2*i*theta) phi_0.  This program verifies:

1. coordinate invariance of the trajectory condition and natural-length ODE;
2. the phase rotation and local prong formula at zeros and a simple pole;
3. the asymptotic-direction count at poles of order at least three;
4. the radial, circular, and logarithmic-spiral cases at a double pole;
5. the Stokes and equal-magnitude rays, dominance sign, and infinity count
   for the Airy equation;
6. the pole order d+4 and d+2 asymptotic directions at infinity for a
   polynomial potential of degree d; and
7. the open action, wall phase, and horizontal/vertical character of the
   real interval for the Weber equation.

Every check raises an explicit exception, so Python's -O flag does not
disable the audit.  The script checks the classical quadratic-differential
geometry; it does not prove Borel summability or a connection formula.
"""

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 failure if condition is false."""

    if not condition:
        raise RuntimeError(message)


def require_zero(expression: sp.Expr, message: str) -> None:
    """Require a symbolic expression to vanish identically."""

    residual = sp.trigsimp(sp.simplify(sp.expand_complex(expression)))
    if residual != 0:
        residual = sp.simplify(residual.rewrite(sp.sin))
    if residual != 0:
        raise RuntimeError(f"{message}: residual = {residual}")


def require_close(
    actual: complex,
    expected: complex,
    tolerance: mp.mpf,
    message: str,
) -> None:
    """Compare complex values using a relative error 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 unit(angle: mp.mpf) -> mp.mpc:
    """Return the unit complex number with the given argument."""

    return mp.exp(1j * angle)


def require_distinct_directions(
    angles: list[mp.mpf],
    message: str,
) -> None:
    """Require all unoriented base rays in a list to be distinct."""

    tolerance = mp.mpf("1e-45")
    directions = [unit(angle) for angle in angles]
    for first in range(len(directions)):
        for second in range(first + 1, len(directions)):
            require(
                abs(directions[first] - directions[second]) > tolerance,
                f"{message}: directions {first} and {second} coincide",
            )


def coordinate_invariance_audit() -> None:
    """Check that the phased trajectory condition survives z=z(u)."""

    theta = sp.symbols("theta", real=True)
    q, jacobian, u_dot, y = sp.symbols(
        "q jacobian u_dot y",
        nonzero=True,
    )

    z_dot = jacobian * u_dot
    base_tangent_value = sp.exp(-2 * sp.I * theta) * q * z_dot**2
    transformed_q = q * jacobian**2
    chart_tangent_value = (
        sp.exp(-2 * sp.I * theta) * transformed_q * u_dot**2
    )
    require_zero(
        chart_tangent_value - base_tangent_value,
        "coordinate invariance of the horizontal-trajectory condition",
    )

    chart_velocity = sp.exp(sp.I * theta) / (y * jacobian)
    pushed_velocity = jacobian * chart_velocity
    require_zero(
        pushed_velocity - sp.exp(sp.I * theta) / y,
        "coordinate covariance of the natural-length trajectory ODE",
    )

    # An explicit branched-coordinate calibration: z=u**2 for Airy.
    u = sp.symbols("u", positive=True)
    z_of_u = u**2
    airy_q_in_u = z_of_u * sp.diff(z_of_u, u) ** 2
    airy_action_in_u = sp.integrate(sp.sqrt(airy_q_in_u), u)
    require_zero(
        airy_q_in_u - 4 * u**4,
        "Airy quadratic differential under z=u**2",
    )
    require_zero(
        airy_action_in_u - sp.Rational(2, 3) * u**3,
        "Airy action under z=u**2",
    )


def phase_and_local_prong_audit() -> None:
    """Verify horizontal and vertical direction formulas locally."""

    theta_value = mp.mpf("0.371")
    coefficient_argument = mp.mpf("0.823")
    coefficient = mp.mpf("1.7") * unit(coefficient_argument)
    tolerance = mp.mpf("1e-45")

    # A zero of order m has m+2 prongs.  The same formula at m=-1
    # gives the single prong at a simple pole.
    for order in (-1, 1, 2, 3, 4):
        count = order + 2
        horizontal_angles = [
            (
                2 * theta_value
                - coefficient_argument
                + 2 * mp.pi * index
            )
            / count
            for index in range(count)
        ]
        vertical_angles = [
            (
                2 * theta_value
                + mp.pi
                - coefficient_argument
                + 2 * mp.pi * index
            )
            / count
            for index in range(count)
        ]
        require_distinct_directions(
            horizontal_angles,
            f"local order {order} horizontal prongs",
        )

        for index, angle in enumerate(horizontal_angles):
            tangent_value = (
                mp.exp(-2j * theta_value)
                * coefficient
                * mp.exp(1j * (order + 2) * angle)
            )
            require_close(
                tangent_value.imag,
                0,
                tolerance,
                f"order {order} horizontal phase at prong {index}",
            )
            require(
                tangent_value.real > 0,
                f"order {order} horizontal prong {index} is not positive",
            )

        for index, angle in enumerate(vertical_angles):
            tangent_value = (
                mp.exp(-2j * theta_value)
                * coefficient
                * mp.exp(1j * (order + 2) * angle)
            )
            require_close(
                tangent_value.imag,
                0,
                tolerance,
                f"order {order} vertical phase at prong {index}",
            )
            require(
                tangent_value.real < 0,
                f"order {order} vertical prong {index} is not negative",
            )

    # At a pole of order n >= 3, q times the squared radial tangent
    # has phase arg(a)-2*theta+(2-n)*arg(x).
    for pole_order in range(3, 8):
        count = pole_order - 2
        angles = [
            (
                coefficient_argument
                - 2 * theta_value
                + 2 * mp.pi * index
            )
            / count
            for index in range(count)
        ]
        require_distinct_directions(
            angles,
            f"pole order {pole_order} asymptotic directions",
        )
        for index, angle in enumerate(angles):
            tangent_value = (
                mp.exp(-2j * theta_value)
                * coefficient
                * mp.exp(1j * (2 - pole_order) * angle)
            )
            require_close(
                tangent_value.imag,
                0,
                tolerance,
                (
                    f"pole order {pole_order} horizontal phase "
                    f"at direction {index}"
                ),
            )
            require(
                tangent_value.real > 0,
                (
                    f"pole order {pole_order} direction {index} "
                    "is not horizontal"
                ),
            )

    # Symbolically check that the proposed natural coordinate differentiates
    # back to a square root of the local quadratic differential.
    x, a = sp.symbols("x a", positive=True)
    for order in (-4, -3, -1, 1, 2, 3, 4):
        natural_coordinate = (
            2
            * sp.sqrt(a)
            * x ** sp.Rational(order + 2, 2)
            / (order + 2)
        )
        require_zero(
            sp.diff(natural_coordinate, x) ** 2 - a * x**order,
            f"local natural coordinate at order {order}",
        )

    graph_period = (
        sp.exp(-2 * sp.I * (sp.symbols("t", real=True) + sp.pi))
        - sp.exp(-2 * sp.I * sp.symbols("t", real=True))
    )
    require_zero(graph_period, "pi-periodicity of the phased graph")


def double_pole_audit() -> None:
    """Check the local double-pole foliation classification."""

    u, v, log_r, alpha = sp.symbols(
        "u v log_r alpha",
        real=True,
    )
    rho = u + sp.I * v
    logarithm = log_r + sp.I * alpha
    imaginary_part = sp.im(sp.expand(rho * logarithm))
    require_zero(
        imaginary_part - (v * log_r + u * alpha),
        "double-pole horizontal-leaf equation",
    )

    constant = mp.mpf("0.7")
    radii = [mp.mpf("0.3"), mp.mpf("0.8"), mp.mpf("2.0")]

    # rho real: alpha is constant while the radius varies.
    radial_u = mp.mpf("2.0")
    radial_angles = [constant / radial_u for _ in radii]
    for angle_value in radial_angles:
        require_close(
            radial_u * angle_value,
            constant,
            mp.mpf("1e-60"),
            "radial double-pole leaf",
        )

    # rho purely imaginary: log(r) is constant while alpha varies.
    circular_v = mp.mpf("1.5")
    circular_radius = mp.exp(constant / circular_v)
    for angle_value in (mp.mpf("-1"), mp.mpf("0"), mp.mpf("1")):
        leaf_value = circular_v * mp.log(circular_radius)
        require_close(
            leaf_value,
            constant,
            mp.mpf("1e-60"),
            f"circular double-pole leaf at alpha={angle_value}",
        )

    # Generic rho: alpha changes linearly with log(r), a logarithmic spiral.
    spiral_u = mp.mpf("2.0")
    spiral_v = mp.mpf("0.5")
    spiral_angles = []
    for radius in radii:
        angle_value = (
            constant - spiral_v * mp.log(radius)
        ) / spiral_u
        spiral_angles.append(angle_value)
        leaf_value = (
            spiral_v * mp.log(radius)
            + spiral_u * angle_value
        )
        require_close(
            leaf_value,
            constant,
            mp.mpf("1e-60"),
            "logarithmic-spiral double-pole leaf",
        )
    require(
        max(spiral_angles) - min(spiral_angles) > mp.mpf("0.1"),
        "generic double-pole leaf did not spiral",
    )

    # The quadratic-differential distance is finite exactly for order > -2.
    epsilon = sp.symbols("epsilon", positive=True)
    radius = sp.symbols("radius", positive=True)
    finite_simple_pole = sp.limit(
        sp.integrate(radius ** sp.Rational(-1, 2), (radius, epsilon, 1)),
        epsilon,
        0,
        dir="+",
    )
    require_zero(
        finite_simple_pole - 2,
        "finite distance to a simple pole",
    )
    double_pole_distance = sp.limit(
        sp.integrate(radius**-1, (radius, epsilon, 1)),
        epsilon,
        0,
        dir="+",
    )
    require(
        double_pole_distance is sp.oo,
        "double pole should be at infinite quadratic-differential distance",
    )


def airy_and_infinity_audit() -> None:
    """Check the Airy rays and polynomial behavior at infinity."""

    theta_value = mp.mpf("0.371")
    tolerance = mp.mpf("1e-45")
    stokes_angles = [
        2 * theta_value / 3 + 2 * mp.pi * index / 3
        for index in range(3)
    ]
    equal_magnitude_angles = [
        (2 * theta_value + (2 * index + 1) * mp.pi) / 3
        for index in range(3)
    ]

    require_distinct_directions(stokes_angles, "Airy Stokes rays")
    require_distinct_directions(
        equal_magnitude_angles,
        "Airy equal-magnitude rays",
    )
    for index, angle in enumerate(stokes_angles):
        rotated_action_phase = unit(-theta_value + 3 * angle / 2)
        require_close(
            rotated_action_phase.real,
            (-1) ** index,
            tolerance,
            f"Airy Stokes action phase at ray {index}",
        )
        require_close(
            rotated_action_phase.imag,
            0,
            tolerance,
            f"Airy Stokes action imaginary part at ray {index}",
        )
    for index, angle in enumerate(equal_magnitude_angles):
        rotated_action_phase = unit(-theta_value + 3 * angle / 2)
        require_close(
            rotated_action_phase.real,
            0,
            tolerance,
            f"Airy equal-magnitude real part at ray {index}",
        )

    # At theta=0 and z>0, W=2*z**(3/2)/3 is positive, so psi_+
    # dominates and psi_- is subdominant.
    hbar_radius = mp.mpf("0.2")
    z_value = mp.mpf("1.3")
    airy_action = mp.mpf(2) * z_value ** mp.mpf("1.5") / 3
    log_magnitude_ratio = 2 * airy_action / hbar_radius
    require(
        log_magnitude_ratio > 0,
        "Airy dominance sign on the positive real ray is reversed",
    )

    # Verify Re(W/hbar)=Re(exp(-i*vartheta)W)/|hbar|.
    action_real, action_imag, vartheta, hbar_abs = sp.symbols(
        "action_real action_imag vartheta hbar_abs",
        real=True,
        positive=True,
    )
    action = action_real + sp.I * action_imag
    hbar = hbar_abs * sp.exp(sp.I * vartheta)
    require_zero(
        sp.re(action / hbar)
        - sp.re(sp.exp(-sp.I * vartheta) * action) / hbar_abs,
        "WKB magnitude phase identity",
    )

    # A degree-d polynomial coefficient becomes a pole of order d+4 at
    # infinity because dz**2 contributes x**(-4), x=1/z.
    x = sp.symbols("x", positive=True)
    for degree in range(0, 7):
        transformed_coefficient = (
            (1 / x) ** degree
            * sp.diff(1 / x, x) ** 2
        )
        require_zero(
            transformed_coefficient - x ** (-(degree + 4)),
            f"degree {degree} polynomial order at infinity",
        )

        pole_order = degree + 4
        direction_count = pole_order - 2
        infinity_angles = [
            (-2 * theta_value + 2 * mp.pi * index) / direction_count
            for index in range(direction_count)
        ]
        require_distinct_directions(
            infinity_angles,
            f"degree {degree} polynomial infinity directions",
        )
        for index, angle in enumerate(infinity_angles):
            tangent_value = (
                mp.exp(-2j * theta_value)
                * mp.exp(1j * (2 - pole_order) * angle)
            )
            require_close(
                tangent_value.imag,
                0,
                tolerance,
                f"degree {degree} infinity direction {index} is not real",
            )
            require(
                tangent_value.real > 0,
                f"degree {degree} infinity direction {index} is not positive",
            )

    # For Airy, map each z-ray to its x=1/z direction and compare with
    # the three directions obtained directly from the order-five pole.
    mapped_to_infinity = [unit(-angle) for angle in stokes_angles]
    pole_directions = [
        unit((-2 * theta_value + 2 * mp.pi * index) / 3)
        for index in range(3)
    ]
    for index, direction in enumerate(mapped_to_infinity):
        nearest = min(abs(direction - candidate) for candidate in pole_directions)
        require(
            nearest < tolerance,
            f"Airy ray {index} does not match an infinity direction",
        )


def weber_audit() -> None:
    """Check the Weber saddle action and its critical phase."""

    a = sp.symbols("a", positive=True, real=True)
    t = sp.symbols("t", real=True)
    semicircle_area = sp.integrate(
        a**2 * sp.cos(t) ** 2,
        (t, -sp.pi / 2, sp.pi / 2),
    )
    require_zero(
        semicircle_area - sp.pi * a**2 / 2,
        "Weber classically allowed action magnitude",
    )

    open_action = sp.I * semicircle_area
    closed_period = 2 * open_action
    require_zero(
        open_action - sp.I * sp.pi * a**2 / 2,
        "Weber open action",
    )
    require_zero(
        closed_period - sp.I * sp.pi * a**2,
        "Weber closed period",
    )

    wall_rotated_action = sp.exp(-sp.I * sp.pi / 2) * open_action
    require_zero(
        wall_rotated_action - sp.pi * a**2 / 2,
        "Weber wall phase theta=pi/2",
    )

    # The parametrization z=a*sin(t) covers the entire open interval
    # between the turning points when |t|<pi/2.
    z_on_interval = a * sp.sin(t)
    q_on_interval = sp.expand(z_on_interval**2 - a**2)
    require_zero(
        q_on_interval + a**2 * sp.cos(t) ** 2,
        "Weber coefficient along the turning-point interval",
    )

    # Writing t=atan(u) makes the strict sign on that open interval a
    # symbolic consequence of a>0 and 1+u**2>0.
    u = sp.symbols("u", real=True)
    q_on_open_interval = sp.simplify(q_on_interval.subs(t, sp.atan(u)))
    require_zero(
        q_on_open_interval + a**2 / (1 + u**2),
        "Weber open-interval coefficient",
    )
    require(
        q_on_open_interval.is_negative is True,
        "Weber interval is not vertical at theta=0",
    )

    wall_q_on_open_interval = sp.simplify(
        sp.exp(-sp.I * sp.pi) * q_on_open_interval
    )
    require_zero(
        wall_q_on_open_interval - a**2 / (1 + u**2),
        "Weber wall-phase coefficient on the open interval",
    )
    require(
        wall_q_on_open_interval.is_positive is True,
        "Weber interval is not horizontal at theta=pi/2",
    )


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

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

    coordinate_invariance_audit()
    print("PASS coordinate invariance and natural-length trajectory ODE")

    phase_and_local_prong_audit()
    print("PASS phased local prongs and higher-pole directions")

    double_pole_audit()
    print("PASS double-pole radial, circular, and spiral leaves")
    print("PASS finite/infinite quadratic-differential distance split")

    airy_and_infinity_audit()
    print("PASS Airy Stokes, equal-magnitude, and dominance rays")
    print("PASS polynomial infinity orders and asymptotic directions")

    weber_audit()
    print("PASS Weber action, wall phase, and interval character")

    print("All Stokes-graph geometry 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)