#!/usr/bin/env python3
"""Audit Razavy determinant zeros with a recessive endpoint connection.

The operator is

    H = -d^2/dx^2 + (zeta*cosh(2*x) - M)^2  on L^2(R).

At the right irregular end, the displayed asymptotic selects the recessive
solution line used by an exact-WKB construction.  We initialize that branch
at a finite cutoff and propagate the *exact* ODE projectively through a
Pruefer angle.  The zeros of the even and odd parity connection coefficients
are then found from the unwrapped angle at x=0.

Because the amplitude is discarded, this locates determinant zeros but does
not evaluate canonically normalized determinant values.  It is also not a
Borel–Padé computation of high-order closed quantum periods.  SciPy >= 1.9
and NumPy are required.
"""

from __future__ import annotations

import argparse
import math
import platform
import sys
from dataclasses import dataclass

import numpy as np

try:
    import scipy
    from scipy.integrate import solve_ivp
    from scipy.optimize import brentq
except ImportError as exc:  # pragma: no cover - user-facing dependency guard
    raise SystemExit(
        "razavy-wkb-connection.py requires SciPy and NumPy"
    ) from exc


@dataclass(frozen=True)
class Profile:
    x_right: float
    x_check: float
    rtol: float
    atol: float
    max_step: float


@dataclass(frozen=True)
class Level:
    index: int
    energy: float
    cutoff_energy: float
    angle_residual: float


PROFILES = {
    "default": Profile(3.0, 2.25, 2.0e-11, 2.0e-12, 0.012),
    "high": Profile(3.5, 2.75, 2.0e-12, 2.0e-13, 0.008),
}


def require(condition: bool, message: str) -> None:
    if not condition:
        raise ValueError(message)


def potential(x: float, zeta: float, parameter_m: float) -> float:
    return (zeta * math.cosh(2.0 * x) - parameter_m) ** 2


def right_log_derivative(
    x: float, energy: float, zeta: float, parameter_m: float
) -> float:
    """Two-term logarithmic derivative of the right-recessive branch."""

    first_correction = -(
        (zeta * zeta + 2.0 * parameter_m - 1.0 - energy) / zeta
    ) * math.exp(-2.0 * x)
    return (
        parameter_m
        - 1.0
        - zeta * math.sinh(2.0 * x)
        + first_correction
    )


def right_angle(
    x: float, energy: float, zeta: float, parameter_m: float
) -> float:
    """Return theta for psi=R sin(theta), psi'=R cos(theta)."""

    angle = math.atan2(
        1.0, right_log_derivative(x, energy, zeta, parameter_m)
    )
    return angle if angle >= 0.0 else angle + 2.0 * math.pi


def propagate_angle(
    energy: float,
    zeta: float,
    parameter_m: float,
    x_right: float,
    profile: Profile,
) -> tuple[float, int]:
    """Propagate the unwrapped Pruefer angle from x_right to the midpoint."""

    def equation(x: float, theta_array: np.ndarray) -> np.ndarray:
        theta = float(theta_array[0])
        sine = math.sin(theta)
        cosine = math.cos(theta)
        derivative = (
            cosine * cosine
            + (energy - potential(x, zeta, parameter_m)) * sine * sine
        )
        return np.array([derivative], dtype=float)

    solution = solve_ivp(
        equation,
        (x_right, 0.0),
        np.array([right_angle(x_right, energy, zeta, parameter_m)]),
        method="DOP853",
        rtol=profile.rtol,
        atol=profile.atol,
        max_step=profile.max_step,
    )
    require(solution.success, f"angle propagation failed: {solution.message}")
    require(
        math.isfinite(float(solution.y[0, -1])),
        "angle propagation returned a non-finite value",
    )
    return float(solution.y[0, -1]), int(solution.nfev)


def target_angle(index: int) -> float:
    """Parity and node-count target for the chosen unwrapped branch."""

    return 0.5 * (1.0 - index) * math.pi


def residual(
    energy: float,
    index: int,
    zeta: float,
    parameter_m: float,
    x_right: float,
    profile: Profile,
) -> float:
    theta, _ = propagate_angle(
        energy, zeta, parameter_m, x_right, profile
    )
    return theta - target_angle(index)


def minimum_potential(zeta: float, parameter_m: float) -> float:
    if parameter_m >= zeta:
        return 0.0
    return (zeta - parameter_m) ** 2


def bracket_level(
    index: int,
    previous_energy: float | None,
    zeta: float,
    parameter_m: float,
    x_right: float,
    profile: Profile,
) -> tuple[float, float]:
    """Bracket one angle target without using the QES polynomial."""

    if previous_energy is None:
        lower = minimum_potential(zeta, parameter_m)
    else:
        lower = previous_energy + 2.0e-8 * max(1.0, abs(previous_energy))

    lower_value = residual(
        lower, index, zeta, parameter_m, x_right, profile
    )
    require(lower_value > 0.0, "lower endpoint lies beyond target angle")

    step = max(1.0, 0.12 * (1.0 + abs(lower)))
    upper = lower + step
    for _ in range(80):
        upper_value = residual(
            upper, index, zeta, parameter_m, x_right, profile
        )
        if upper_value < 0.0:
            return lower, upper
        step *= 1.45
        upper += step
    raise ValueError(f"could not bracket level {index}")


def solve_levels(
    count: int,
    zeta: float,
    parameter_m: float,
    x_right: float,
    profile: Profile,
) -> list[float]:
    energies: list[float] = []
    for index in range(count):
        lower, upper = bracket_level(
            index,
            energies[-1] if energies else None,
            zeta,
            parameter_m,
            x_right,
            profile,
        )
        root = brentq(
            residual,
            lower,
            upper,
            args=(index, zeta, parameter_m, x_right, profile),
            xtol=5.0e-13,
            rtol=2.0e-14,
            maxiter=120,
        )
        energies.append(float(root))
    return energies


def qes_roots(zeta: float, parameter_m: float) -> list[float] | None:
    """Return the exact M=4 roots used by the chapter calibration."""

    if abs(zeta - 1.0) > 1.0e-14 or abs(parameter_m - 4.0) > 1.0e-14:
        return None
    return sorted([6.0, 14.0 - 4.0 * math.sqrt(3.0), 14.0,
                   14.0 + 4.0 * math.sqrt(3.0)])


def parse_arguments() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--profile", choices=tuple(PROFILES), default="default")
    parser.add_argument("--zeta", type=float, default=1.0)
    parser.add_argument("--M", dest="parameter_m", type=float, default=4.0)
    parser.add_argument("--levels", type=int, default=5)
    return parser.parse_args()


def main() -> int:
    arguments = parse_arguments()
    require(math.isfinite(arguments.zeta) and arguments.zeta > 0.0,
            "--zeta must be positive and finite")
    require(
        math.isfinite(arguments.parameter_m) and arguments.parameter_m > 0.0,
        "--M must be positive and finite",
    )
    require(1 <= arguments.levels <= 10, "--levels must lie between 1 and 10")

    profile = PROFILES[arguments.profile]
    final_energies = solve_levels(
        arguments.levels,
        arguments.zeta,
        arguments.parameter_m,
        profile.x_right,
        profile,
    )
    cutoff_energies = solve_levels(
        arguments.levels,
        arguments.zeta,
        arguments.parameter_m,
        profile.x_check,
        profile,
    )

    rows: list[Level] = []
    for index, energy in enumerate(final_energies):
        rows.append(
            Level(
                index=index,
                energy=energy,
                cutoff_energy=cutoff_energies[index],
                angle_residual=abs(
                    residual(
                        energy,
                        index,
                        arguments.zeta,
                        arguments.parameter_m,
                        profile.x_right,
                        profile,
                    )
                ),
            )
        )

    exact = qes_roots(arguments.zeta, arguments.parameter_m)
    exact_gaps = [] if exact is None else [
        abs(rows[index].energy - exact[index])
        for index in range(min(len(exact), len(rows)))
    ]
    cutoff_shifts = [abs(row.energy - row.cutoff_energy) for row in rows]
    maximum_angle_residual = max(row.angle_residual for row in rows)

    print("Razavy sectorial-WKB connection audit")
    print(f"Python {platform.python_version()}; NumPy {np.__version__}; "
          f"SciPy {scipy.__version__}")
    print(
        f"profile={arguments.profile}; zeta={arguments.zeta:g}; "
        f"M={arguments.parameter_m:g}; x_R={profile.x_right:g}; "
        f"cutoff check={profile.x_check:g}"
    )
    print("right tail: exp[-zeta*cosh(2x)/2 + (M-1)x]")
    print("target: theta(0; E_n)=(1-n)*pi/2")
    print()
    print(" n parity      QES energy       connection root    cutoff shift   angle residual")
    for row in rows:
        qes = "--" if exact is None or row.index >= len(exact) else f"{exact[row.index]:.12f}"
        parity = "even" if row.index % 2 == 0 else "odd"
        print(
            f"{row.index:2d} {parity:>5s}  {qes:>16s}  "
            f"{row.energy:20.12f}  "
            f"{abs(row.energy-row.cutoff_energy):12.3e}  "
            f"{row.angle_residual:14.3e}"
        )

    print()
    if exact_gaps:
        print(f"maximum QES/connection gap: {max(exact_gaps):.3e}")
    print(f"maximum endpoint-cutoff shift: {max(cutoff_shifts):.3e}")
    print(f"maximum angle residual: {maximum_angle_residual:.3e}")
    print(
        "Scope: projective exact-ODE propagation of parity determinant-zero "
        "conditions; no determinant amplitude or Borel–Padé period sum."
    )

    qes_gate = not exact_gaps or max(exact_gaps) <= 2.0e-9
    cutoff_gate = max(cutoff_shifts) <= 2.0e-8
    angle_gate = maximum_angle_residual <= 2.0e-10
    require(qes_gate, "QES/connection regression gate failed")
    require(cutoff_gate, "endpoint-cutoff regression gate failed")
    require(angle_gate, "angle-residual regression gate failed")
    print("Regression gates: PASS")
    return 0


if __name__ == "__main__":
    try:
        sys.exit(main())
    except ValueError as error:
        print(f"error: {error}", file=sys.stderr)
        sys.exit(2)