#!/usr/bin/env python3
"""Audit Chapter 9 Page 8 complex-resonance and PT formulas.

The checks deliberately use independent representations:

1. asymptotic phase geometry for polynomial decay wedges;
2. symbolic Hermite residuals for the outgoing inverted oscillator;
3. a finite-difference complex-scaled oscillator;
4. the exact Airy determinant of the PT-symmetric imaginary-linear box;
5. high-precision exceptional-point equations and Puiseux splitting;
6. an unrelated finite-difference Airy-box discretization.

Every failure raises an explicit exception.  No Python assert statements
are used, so normal and optimized execution perform the same audit.
The script does not prove dilation analyticity, Borel summability,
resolvent continuation, or completeness for a new spectral problem.
"""

from __future__ import annotations

import cmath
import math
import platform

import mpmath as mp
import numpy as np
import scipy
from scipy.linalg import eigvals
import sympy as sp


def require(condition: bool, message: str) -> None:
    """Raise an optimization-safe failure unless condition is true."""

    if not condition:
        raise RuntimeError(message)


def require_close(
    actual: complex,
    expected: complex,
    tolerance: float | mp.mpf,
    message: str,
) -> None:
    """Compare numbers with a scale-aware tolerance."""

    error = abs(actual - expected)
    scale = max(1, abs(actual), abs(expected))
    if error > tolerance * scale:
        raise RuntimeError(
            f"{message}: error {mp.nstr(error, 14)} exceeds "
            f"{mp.nstr(tolerance * scale, 14)}"
        )


def wedge_audit() -> None:
    """Verify the number and spacing of polynomial decay sectors."""

    for degree in range(1, 9):
        count = degree + 2
        arg_a = 0.37
        arg_hbar = -0.21
        rays = []
        for index in range(count):
            angle = (
                2
                / count
                * (
                    (index + 0.5) * math.pi
                    - arg_a / 2
                    + arg_hbar
                )
            )
            rays.append(angle % (2 * math.pi))
        rays.sort()
        spacings = [
            (rays[(j + 1) % count] - rays[j]) % (2 * math.pi)
            for j in range(count)
        ]
        expected = 2 * math.pi / count
        for spacing in spacings:
            require_close(
                spacing,
                expected,
                2e-14,
                f"degree-{degree} wedge spacing",
            )


def inverted_oscillator_symbolic_audit() -> None:
    """Check outgoing Hermite states against the original ODE."""

    z = sp.symbols("z")
    hbar = sp.symbols("hbar", positive=True)
    rotation = sp.exp(-sp.I * sp.pi / 4)

    for level in range(5):
        state = sp.hermite(
            level,
            rotation * z / sp.sqrt(hbar),
        ) * sp.exp(sp.I * z**2 / (2 * hbar))
        residual = (
            -hbar**2 * sp.diff(state, z, 2)
            - z**2 * state
            + sp.I * hbar * (2 * level + 1) * state
        )
        reduced = sp.simplify(sp.expand(residual))
        require(
            reduced == 0,
            f"Hermite resonance residual at level {level}: {reduced}",
        )

    theta = math.pi / 4
    kinetic_phase = cmath.exp(-2j * theta)
    potential_phase = -cmath.exp(2j * theta)
    require_close(
        kinetic_phase,
        -1j,
        2e-15,
        "rotated inverted-oscillator kinetic phase",
    )
    require_close(
        potential_phase,
        -1j,
        2e-15,
        "rotated inverted-oscillator potential phase",
    )


def second_difference_matrix(points: int, half_width: float) -> tuple:
    """Return the Dirichlet grid and matrix for minus the second derivative."""

    step = 2 * half_width / (points + 1)
    grid = np.linspace(
        -half_width + step,
        half_width - step,
        points,
    )
    main = np.full(points, 2 / step**2)
    off = np.full(points - 1, -1 / step**2)
    kinetic = (
        np.diag(main)
        + np.diag(off, 1)
        + np.diag(off, -1)
    )
    return grid, kinetic


def scaled_oscillator_matrix_audit() -> None:
    """Compare a rotated finite-difference matrix with exact poles."""

    grid, kinetic = second_difference_matrix(220, 9.0)
    theta = math.pi / 4
    matrix = (
        cmath.exp(-2j * theta) * kinetic
        - cmath.exp(2j * theta) * np.diag(grid**2)
    )
    spectrum = eigvals(
        matrix,
        overwrite_a=True,
        check_finite=False,
    )

    tolerances = (8e-4, 3e-3, 7e-3)
    for level, tolerance in zip(range(3), tolerances):
        expected = -1j * (2 * level + 1)
        error = float(np.min(np.abs(spectrum - expected)))
        require(
            error < tolerance,
            f"scaled oscillator level {level}: error {error}",
        )


def airy_delta(energy: complex, coupling: complex) -> complex:
    """Return the real-normalized Airy-box spectral determinant."""

    alpha = coupling ** (mp.mpf(1) / 3) * mp.exp(mp.j * mp.pi / 6)
    z_minus = alpha * (-1 + mp.j * energy / coupling)
    z_plus = alpha * (1 + mp.j * energy / coupling)
    raw = (
        mp.airyai(z_minus) * mp.airybi(z_plus)
        - mp.airyai(z_plus) * mp.airybi(z_minus)
    )
    return mp.exp(-mp.j * mp.pi / 6) * raw


def airy_reality_audit() -> None:
    """Check the fixed phase that makes the PT determinant real."""

    mp.mp.dps = 70
    samples = (
        (mp.mpf("1.7"), mp.mpf("1.0")),
        (mp.mpf("7.0"), mp.mpf("10.0")),
        (mp.mpf("15.0"), mp.mpf("20.0")),
    )
    for energy, coupling in samples:
        value = airy_delta(energy, coupling)
        require_close(
            mp.im(value),
            0,
            mp.mpf("1e-62"),
            "PT Airy determinant reality phase",
        )


def find_exceptional_point() -> tuple:
    """Solve the Airy determinant and derivative equations."""

    mp.mp.dps = 70

    def derivative(energy: complex, coupling: complex) -> complex:
        return mp.diff(
            lambda variable: airy_delta(variable, coupling),
            energy,
        )

    energy, coupling = mp.findroot(
        lambda e, g: (
            mp.re(airy_delta(e, g)),
            mp.re(derivative(e, g)),
        ),
        (mp.mpf("7.1"), mp.mpf("12.3")),
        tol=mp.mpf("1e-60"),
        maxsteps=80,
    )
    require_close(
        airy_delta(energy, coupling),
        0,
        mp.mpf("1e-55"),
        "Airy exceptional-point determinant",
    )
    require_close(
        derivative(energy, coupling),
        0,
        mp.mpf("1e-52"),
        "Airy exceptional-point energy derivative",
    )
    require_close(
        energy,
        mp.mpf(
            "7.1085995967649487972434189646670883557178474563231"
        ),
        mp.mpf("1e-48"),
        "Airy exceptional energy",
    )
    require_close(
        coupling,
        mp.mpf(
            "12.312455672260525052957083672793968699368624481258"
        ),
        mp.mpf("1e-48"),
        "Airy exceptional coupling",
    )
    require_close(
        energy / coupling,
        1 / mp.sqrt(3),
        mp.mpf("1e-48"),
        "Airy exceptional-point ratio",
    )
    return energy, coupling


def airy_puiseux_audit(energy: complex, coupling: complex) -> None:
    """Check the local real-to-conjugate square-root splitting."""

    derivative_g = mp.diff(
        lambda variable: airy_delta(energy, variable),
        coupling,
    )
    derivative_ee = mp.diff(
        lambda variable: airy_delta(variable, coupling),
        energy,
        2,
    )
    coefficient_squared = mp.re(-2 * derivative_g / derivative_ee)
    require_close(
        coefficient_squared,
        mp.mpf(
            "-2.388059376232862147971674530659651209971989829569"
        ),
        mp.mpf("1e-46"),
        "Airy Puiseux squared coefficient",
    )
    coefficient = mp.sqrt(-coefficient_squared)

    epsilon = mp.mpf("1e-6")
    displacement = coefficient * mp.sqrt(epsilon)

    coupling_below = coupling - epsilon
    root_minus = mp.findroot(
        lambda variable: airy_delta(variable, coupling_below),
        (energy - 1.2 * displacement, energy - 0.8 * displacement),
        tol=mp.mpf("1e-54"),
    )
    root_plus = mp.findroot(
        lambda variable: airy_delta(variable, coupling_below),
        (energy + 0.8 * displacement, energy + 1.2 * displacement),
        tol=mp.mpf("1e-54"),
    )
    measured_below = (root_plus - root_minus) / (2 * mp.sqrt(epsilon))
    require_close(
        measured_below,
        coefficient,
        mp.mpf("5e-8"),
        "real-side Airy Puiseux coefficient",
    )

    coupling_above = coupling + epsilon
    imaginary_displacement = mp.j * displacement
    root_upper = mp.findroot(
        lambda variable: airy_delta(variable, coupling_above),
        (
            energy + 0.8 * imaginary_displacement,
            energy + 1.2 * imaginary_displacement,
        ),
        tol=mp.mpf("1e-54"),
    )
    root_lower = mp.findroot(
        lambda variable: airy_delta(variable, coupling_above),
        (
            energy - 0.8 * imaginary_displacement,
            energy - 1.2 * imaginary_displacement,
        ),
        tol=mp.mpf("1e-54"),
    )
    measured_above = (
        (root_upper - root_lower)
        / (2 * mp.j * mp.sqrt(epsilon))
    )
    require_close(
        measured_above,
        coefficient,
        mp.mpf("5e-8"),
        "complex-side Airy Puiseux coefficient",
    )
    require_close(
        root_upper,
        mp.conj(root_lower),
        mp.mpf("1e-52"),
        "Airy roots above the EP form a conjugate pair",
    )


def airy_root(guess: complex, coupling: float) -> complex:
    """Find one exact Airy-determinant zero from a complex guess."""

    start = mp.mpc(guess)
    return mp.findroot(
        lambda variable: airy_delta(variable, coupling),
        (start, start * (1 + mp.mpf("1e-4")) + mp.mpf("1e-5")),
        tol=mp.mpf("1e-52"),
        maxsteps=80,
    )


def airy_box_matrix(coupling: float, points: int = 220) -> np.ndarray:
    """Return finite-difference eigenvalues of the PT Airy box."""

    grid, kinetic = second_difference_matrix(points, 1.0)
    matrix = kinetic.astype(complex) + 1j * coupling * np.diag(grid)
    return eigvals(
        matrix,
        overwrite_a=True,
        check_finite=False,
    )


def airy_independent_matrix_audit() -> None:
    """Compare exact Airy roots with an unrelated matrix discretization."""

    exact_below = (
        airy_root(4.6, 10.0),
        airy_root(9.0, 10.0),
    )
    expected_below = (
        mp.mpf("4.5784904182916919656"),
        mp.mpf("8.9966832557123482084"),
    )
    for actual, expected in zip(exact_below, expected_below):
        require_close(
            actual,
            expected,
            mp.mpf("1e-19"),
            "exact Airy root below the EP",
        )

    exact_upper = airy_root(7.2 + 1.3j, 13.0)
    exact_lower = airy_root(7.2 - 1.3j, 13.0)
    require_close(
        exact_upper,
        mp.mpc(
            "7.2169577471438904527",
            "1.3047650935833622415",
        ),
        mp.mpf("1e-19"),
        "exact Airy root above the EP",
    )
    require_close(
        exact_upper,
        mp.conj(exact_lower),
        mp.mpf("1e-50"),
        "exact Airy conjugate pairing at g=13",
    )

    matrix_below = airy_box_matrix(10.0)
    matrix_above = airy_box_matrix(13.0)
    for target in exact_below:
        error = float(
            np.min(np.abs(matrix_below - complex(target)))
        )
        require(
            error < 8e-4,
            f"finite-difference Airy root below EP: error {error}",
        )
    for target in (exact_upper, exact_lower):
        error = float(
            np.min(np.abs(matrix_above - complex(target)))
        )
        require(
            error < 8e-4,
            f"finite-difference Airy root above EP: error {error}",
        )


def main() -> None:
    """Run the full optimization-safe audit."""

    print("Advanced ODE Chapter 9 Page 8 complex-spectrum audit")
    print(f"Python: {platform.python_version()}")
    print(f"SymPy: {sp.__version__}")
    print(f"mpmath: {mp.__version__}")
    print(f"NumPy: {np.__version__}")
    print(f"SciPy: {scipy.__version__}")

    wedge_audit()
    print("PASS polynomial decay-wedge count and spacing")

    inverted_oscillator_symbolic_audit()
    print("PASS inverted-oscillator rotation and Hermite resonances")

    scaled_oscillator_matrix_audit()
    print("PASS finite-difference complex-scaled oscillator")

    airy_reality_audit()
    print("PASS PT-normalized Airy determinant reality")

    energy, coupling = find_exceptional_point()
    print("PASS high-precision Airy determinant double zero")

    airy_puiseux_audit(energy, coupling)
    print("PASS Airy exceptional-point Puiseux splitting")

    airy_independent_matrix_audit()
    print("PASS exact Airy roots and independent matrix spectrum")

    print(
        "LIMITATION: these identities do not prove dilation analyticity, "
        "resolvent continuation, Borel summability, or spectral "
        "completeness for a new problem."
    )
    print("All complex-resonance and PT checks passed.")


if __name__ == "__main__":
    main()