#!/usr/bin/env python3
"""Audit the spectral branch through the one-state Razavy QES point.

The companion page uses

    H(M) = -d^2/dx^2 + (zeta*cosh(2*x) - M)^2,
    zeta > 0.

At M=1,

    E_0 = zeta^2 + 1,
    psi_0 is proportional to exp[-zeta*cosh(2*x)/2],

and Hellmann--Feynman gives

    E_0'(1) = 2*(1 - zeta*K_1(zeta)/K_0(zeta)).

This script evaluates the Bessel-integral ratio by Gauss--Legendre
quadrature and compares the exact slope with a centered finite-difference
of independently refined coordinate-grid eigenvalues at M=1+/-delta.
It also writes the ground-state curve used in the companion figure.

The half-line grid is the even block of a centered second-order Dirichlet
grid.  Each tridiagonal ground eigenvalue is found by Sturm bisection; step
halving and second-order Richardson extrapolation remove the leading mesh
error.  Reported errors are empirical diagnostics, not certified bounds.

Requirements: CPython 3.9+ and NumPy 1.24+.
"""

from __future__ import annotations

import argparse
import csv
from dataclasses import dataclass
import math
from pathlib import Path
from typing import Sequence

import numpy as np


@dataclass(frozen=True)
class Settings:
    """Numerical settings for one refinement profile."""

    coarse_cells: int
    domain: float
    sturm_iterations: int
    quadrature_nodes: int


@dataclass(frozen=True)
class RefinedEnergy:
    """One Richardson-refined ground-state energy."""

    value: float
    fine_value: float
    refinement_shift: float
    sturm_width: float


def require(condition: bool, message: str) -> None:
    """Raise an optimization-safe exception when an audit condition fails."""

    if not condition:
        raise RuntimeError(message)


def settings_for_profile(profile: str) -> Settings:
    """Return calibrated default or high-refinement settings."""

    if profile == "default":
        return Settings(
            coarse_cells=1800,
            domain=3.5,
            sturm_iterations=88,
            quadrature_nodes=256,
        )
    if profile == "high":
        return Settings(
            coarse_cells=3600,
            domain=3.5,
            sturm_iterations=96,
            quadrature_nodes=384,
        )
    raise ValueError(f"unknown profile {profile!r}")


def potential(x: np.ndarray, zeta: float, m_value: float) -> np.ndarray:
    """Return the Razavy potential for real, not necessarily integer, M."""

    return (zeta * np.cosh(2.0 * x) - m_value) ** 2


def even_grid_block(
    zeta: float,
    m_value: float,
    domain: float,
    cells: int,
) -> tuple[np.ndarray, np.ndarray]:
    """Return the parity-even tridiagonal block on the half line."""

    require(domain >= 3.0, "domain must be at least 3")
    require(cells >= 600, "coarse grid must have at least 600 cells")
    step = domain / cells
    diagonal = (
        2.0 / step**2
        + potential(step * np.arange(cells, dtype=float), zeta, m_value)
    )
    off_diagonal = np.full(cells - 1, -1.0 / step**2)
    # Orthogonal parity reduction of the centered full-line grid.
    off_diagonal[0] *= math.sqrt(2.0)
    return np.asarray(diagonal), off_diagonal


def sturm_count(
    diagonal: np.ndarray,
    off_diagonal: np.ndarray,
    value: float,
) -> int:
    """Count tridiagonal eigenvalues strictly below value."""

    pivot_floor = 8.0 * np.finfo(float).eps * max(1.0, abs(value))
    pivot = float(diagonal[0] - value)
    count = int(pivot < 0.0)
    for index, coupling in enumerate(off_diagonal):
        if abs(pivot) < pivot_floor:
            pivot = -pivot_floor if pivot < 0.0 else pivot_floor
        pivot = float(
            diagonal[index + 1]
            - value
            - coupling * coupling / pivot
        )
        count += int(pivot < 0.0)
    return count


def tridiagonal_bounds(
    diagonal: np.ndarray,
    off_diagonal: np.ndarray,
) -> tuple[float, float]:
    """Return Gershgorin bounds containing the full spectrum."""

    radii = np.zeros_like(diagonal)
    radii[:-1] += np.abs(off_diagonal)
    radii[1:] += np.abs(off_diagonal)
    return (
        float(np.min(diagonal - radii)),
        float(np.max(diagonal + radii)),
    )


def ground_once(
    zeta: float,
    m_value: float,
    domain: float,
    cells: int,
    iterations: int,
) -> tuple[float, float]:
    """Find the lowest even-grid eigenvalue by Sturm bisection."""

    diagonal, off_diagonal = even_grid_block(
        zeta, m_value, domain, cells
    )
    lower, upper = tridiagonal_bounds(diagonal, off_diagonal)
    for _ in range(iterations):
        midpoint = 0.5 * (lower + upper)
        if midpoint == lower or midpoint == upper:
            break
        if sturm_count(diagonal, off_diagonal, midpoint) == 0:
            lower = midpoint
        else:
            upper = midpoint
    return 0.5 * (lower + upper), upper - lower


def refined_ground(
    zeta: float,
    m_value: float,
    settings: Settings,
) -> RefinedEnergy:
    """Step-halve and Richardson-extrapolate the ground energy."""

    coarse, coarse_width = ground_once(
        zeta,
        m_value,
        settings.domain,
        settings.coarse_cells,
        settings.sturm_iterations,
    )
    fine, fine_width = ground_once(
        zeta,
        m_value,
        settings.domain,
        2 * settings.coarse_cells,
        settings.sturm_iterations,
    )
    richardson = (4.0 * fine - coarse) / 3.0
    return RefinedEnergy(
        value=float(richardson),
        fine_value=float(fine),
        refinement_shift=float(abs(richardson - fine)),
        sturm_width=max(coarse_width, fine_width),
    )


def bessel_ratio_k1_k0(zeta: float, nodes: int) -> float:
    """Evaluate K_1(zeta)/K_0(zeta) from its integral representation."""

    # For zeta >= 0.5, the tail beyond t=12 is far below double precision.
    cutoff = 12.0
    raw_nodes, raw_weights = np.polynomial.legendre.leggauss(nodes)
    t_values = 0.5 * cutoff * (raw_nodes + 1.0)
    weights = 0.5 * cutoff * raw_weights
    common = np.exp(-zeta * np.cosh(t_values))
    k_zero = float(np.dot(weights, common))
    k_one = float(np.dot(weights, common * np.cosh(t_values)))
    require(k_zero > 0.0, "K_0 quadrature underflowed")
    return k_one / k_zero


def exact_slope(zeta: float, nodes: int) -> float:
    """Return the exact Hellmann--Feynman slope at M=1."""

    return 2.0 * (1.0 - zeta * bessel_ratio_k1_k0(zeta, nodes))


def finite_float(text: str) -> float:
    """Parse one finite command-line float."""

    value = float(text)
    if not math.isfinite(value):
        raise argparse.ArgumentTypeError("value must be finite")
    return value


def write_csv(
    path: Path,
    zeta: float,
    curve: list[tuple[float, RefinedEnergy]],
    delta: float,
    analytic_slope: float,
    grid_slope: float,
    settings: Settings,
) -> None:
    """Write curve points and provenance in a rectangular CSV."""

    path.parent.mkdir(parents=True, exist_ok=True)
    fieldnames = [
        "row_type",
        "zeta",
        "M",
        "energy",
        "fine_energy",
        "refinement_shift",
        "sturm_width",
        "delta",
        "analytic_slope",
        "grid_slope",
        "frozen_block_slope",
        "coarse_cells",
        "fine_cells",
        "domain",
        "sturm_iterations",
        "quadrature_nodes",
    ]
    with path.open("w", newline="", encoding="utf-8") as handle:
        writer = csv.DictWriter(handle, fieldnames=fieldnames)
        writer.writeheader()
        for m_value, result in curve:
            writer.writerow(
                {
                    "row_type": "curve",
                    "zeta": f"{zeta:.17g}",
                    "M": f"{m_value:.17g}",
                    "energy": f"{result.value:.17g}",
                    "fine_energy": f"{result.fine_value:.17g}",
                    "refinement_shift": f"{result.refinement_shift:.17g}",
                    "sturm_width": f"{result.sturm_width:.17g}",
                    "delta": f"{delta:.17g}",
                    "analytic_slope": f"{analytic_slope:.17g}",
                    "grid_slope": f"{grid_slope:.17g}",
                    "frozen_block_slope": "2",
                    "coarse_cells": settings.coarse_cells,
                    "fine_cells": 2 * settings.coarse_cells,
                    "domain": f"{settings.domain:.17g}",
                    "sturm_iterations": settings.sturm_iterations,
                    "quadrature_nodes": settings.quadrature_nodes,
                }
            )


def parse_arguments(argv: Sequence[str] | None = None) -> argparse.Namespace:
    """Parse and validate the command line."""

    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--zeta",
        type=finite_float,
        default=1.0,
        help="Razavy zeta (calibrated range: 0.5 to 2; default: 1)",
    )
    parser.add_argument(
        "--delta",
        type=finite_float,
        default=0.01,
        help="centered detuning step (range: 0.002 to 0.03; default: 0.01)",
    )
    parser.add_argument(
        "--profile",
        choices=("default", "high"),
        default="default",
        help="grid and quadrature refinement profile",
    )
    parser.add_argument(
        "--scan-min",
        type=finite_float,
        default=0.75,
        help="left endpoint of the plotted M scan",
    )
    parser.add_argument(
        "--scan-max",
        type=finite_float,
        default=1.25,
        help="right endpoint of the plotted M scan",
    )
    parser.add_argument(
        "--scan-samples",
        type=int,
        default=11,
        help="number of uniformly spaced curve samples (default: 11)",
    )
    parser.add_argument(
        "--csv",
        type=Path,
        help="optional path for curve data and numerical provenance",
    )
    args = parser.parse_args(argv)
    if not 0.5 <= args.zeta <= 2.0:
        parser.error("--zeta must lie in the calibrated range [0.5, 2]")
    if not 0.002 <= args.delta <= 0.03:
        parser.error("--delta must lie in the calibrated range [0.002, 0.03]")
    if not 0.5 <= args.scan_min < 1.0 < args.scan_max <= 1.5:
        parser.error(
            "the calibrated scan interval must contain M=1 and lie in [0.5, 1.5]"
        )
    if not 3 <= args.scan_samples <= 101:
        parser.error("--scan-samples must lie between 3 and 101")
    return args


def main(argv: Sequence[str] | None = None) -> int:
    """Run the analytic, grid, and curve-data audit."""

    args = parse_arguments(argv)
    settings = settings_for_profile(args.profile)
    analytic = exact_slope(args.zeta, settings.quadrature_nodes)

    minus = refined_ground(args.zeta, 1.0 - args.delta, settings)
    center = refined_ground(args.zeta, 1.0, settings)
    plus = refined_ground(args.zeta, 1.0 + args.delta, settings)
    grid_slope = (plus.value - minus.value) / (2.0 * args.delta)
    exact_energy = args.zeta**2 + 1.0
    center_gap = abs(center.value - exact_energy)
    slope_gap = abs(grid_slope - analytic)

    scan_values = np.linspace(
        args.scan_min,
        args.scan_max,
        args.scan_samples,
    )
    curve = [
        (float(m_value), refined_ground(args.zeta, float(m_value), settings))
        for m_value in scan_values
    ]

    scalar_audit = np.asarray(
        [
            analytic,
            grid_slope,
            exact_energy,
            center.value,
            center_gap,
            slope_gap,
        ],
        dtype=float,
    )
    curve_audit = np.asarray(
        [
            [
                m_value,
                result.value,
                result.fine_value,
                result.refinement_shift,
                result.sturm_width,
            ]
            for m_value, result in curve
        ],
        dtype=float,
    )
    require(np.all(np.isfinite(scalar_audit)), "non-finite central audit value")
    require(np.all(np.isfinite(curve_audit)), "non-finite curve audit value")
    require(center_gap < 2.0e-7, "QES/grid energy check failed")
    # The centered derivative has an O(delta^2) error in addition to mesh
    # error.  This calibrated gate remains conservative over the CLI range.
    require(
        slope_gap < 0.20 * args.delta**2 + 2.0e-6,
        "analytic/grid slope check failed",
    )

    print("Razavy QES detuning audit")
    print(f"profile                 = {args.profile}")
    print(f"zeta                    = {args.zeta:.12g}")
    print(f"delta                   = {args.delta:.12g}")
    print(f"exact E0(1)             = {exact_energy:.12f}")
    print(f"grid E0(1)              = {center.value:.12f}")
    print(f"|grid - exact|          = {center_gap:.3e}")
    print(f"exact E0'(1)            = {analytic:.12f}")
    print(f"centered grid slope     = {grid_slope:.12f}")
    print(f"|slope difference|      = {slope_gap:.3e}")
    print("frozen-block slope      = 2.000000000000")
    print(
        "maximum refinement shift = "
        f"{max(item.refinement_shift for _, item in curve):.3e}"
    )
    print(f"maximum Sturm width     = {max(item.sturm_width for _, item in curve):.3e}")
    print("checks                   = PASS")

    if args.csv is not None:
        write_csv(
            args.csv,
            args.zeta,
            curve,
            args.delta,
            analytic,
            grid_slope,
            settings,
        )
        print(f"csv                      = {args.csv}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())