#!/usr/bin/env python3
"""Calibrate several approximations to the stable quartic oscillator.

The dimensionless problem is

    H(g) = -d^2/dx^2 + x^2 + g*x^4,   g >= 0,   psi in L^2(R).

The script compares:

* exact harmonic data and the ground-state Bender--Wu recurrence;
* leading Bohr--Sommerfeld quantization;
* a parity-resolved half-line Fourier--Galerkin calculation; and
* an independent coordinate finite difference with Sturm bisection.

The two numerical methods are refined separately.  Their discrepancies and
refinement shifts are empirical diagnostics, not certified error intervals.

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

from __future__ import annotations

import argparse
import csv
from dataclasses import dataclass
from fractions import Fraction
import math
import platform
import sys
from typing import Iterable

import numpy as np


PI = math.pi


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

    if not condition:
        raise RuntimeError(message)


def box_length(g: float, scale: float = 8.0) -> float:
    """Return a coupling-adapted half-box length."""

    require(g >= 0.0, "the stable calibration requires g >= 0")
    return scale * (1.0 + g) ** (-1.0 / 6.0)


def cosine_moment(power: int, modes: np.ndarray, length: float) -> np.ndarray:
    """Return L^(-1) integral_0^L x^power cos(m*pi*x/L) dx.

    Only powers two and four are needed.  ``modes`` may contain negative
    integers because the answer is even in m.
    """

    m = np.abs(np.asarray(modes, dtype=np.int64))
    answer = np.empty(m.shape, dtype=float)
    zero = m == 0
    nonzero = ~zero
    if power == 2:
        answer[zero] = length**2 / 3.0
    elif power == 4:
        answer[zero] = length**4 / 5.0
    else:
        raise ValueError("only powers two and four are implemented")

    if np.any(nonzero):
        mn = m[nonzero].astype(float)
        sign = np.where(m[nonzero] % 2 == 0, 1.0, -1.0)
        z = mn * PI
        if power == 2:
            answer[nonzero] = 2.0 * length**2 * sign / z**2
        else:
            answer[nonzero] = length**4 * sign * (4.0 / z**2 - 24.0 / z**4)
    return answer


@dataclass
class GalerkinLevel:
    """One Fourier--Galerkin eigenpair diagnostic."""

    energy: float
    parity: str
    local_index: int
    kinetic: float
    x2: float
    x4: float


def fourier_block(
    *,
    parity: str,
    quadratic: float,
    quartic: float,
    length: float,
    basis_size: int,
) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    """Diagonalize one half-line Fourier--Galerkin parity block."""

    require(parity in {"even", "odd"}, "parity must be even or odd")
    require(length > 0.0, "length must be positive")
    require(basis_size >= 8, "basis_size is too small for this calibration")

    j = np.arange(basis_size, dtype=np.int64)
    jj, kk = np.meshgrid(j, j, indexing="ij")
    if parity == "even":
        wave_numbers = (j + 0.5) * PI / length
        x2 = cosine_moment(2, jj - kk, length) + cosine_moment(
            2, jj + kk + 1, length
        )
        x4 = cosine_moment(4, jj - kk, length) + cosine_moment(
            4, jj + kk + 1, length
        )
    else:
        wave_numbers = (j + 1.0) * PI / length
        x2 = cosine_moment(2, jj - kk, length) - cosine_moment(
            2, jj + kk + 2, length
        )
        x4 = cosine_moment(4, jj - kk, length) - cosine_moment(
            4, jj + kk + 2, length
        )

    kinetic = np.diag(wave_numbers**2)
    hamiltonian = kinetic + quadratic * x2 + quartic * x4
    energies, vectors = np.linalg.eigh(hamiltonian)
    return energies, vectors, kinetic, x2, x4


def fourier_levels(
    g: float,
    count: int,
    *,
    basis_size: int,
    length: float,
) -> list[GalerkinLevel]:
    """Return the lowest full-line levels from two parity blocks."""

    candidates: list[GalerkinLevel] = []
    per_parity = (count + 1) // 2 + 1
    for parity in ("even", "odd"):
        energies, vectors, kinetic, x2, x4 = fourier_block(
            parity=parity,
            quadratic=1.0,
            quartic=g,
            length=length,
            basis_size=basis_size,
        )
        for local_index in range(per_parity):
            vector = vectors[:, local_index]
            expectation = lambda matrix: float(vector @ (matrix @ vector))
            candidates.append(
                GalerkinLevel(
                    energy=float(energies[local_index]),
                    parity=parity,
                    local_index=local_index,
                    kinetic=expectation(kinetic),
                    x2=expectation(x2),
                    x4=expectation(x4),
                )
            )
    candidates.sort(key=lambda level: level.energy)
    return candidates[:count]


def virial_residual(level: GalerkinLevel, g: float) -> float:
    """Return the quartic virial residual for one approximate eigenvector."""

    return 2.0 * level.kinetic - 2.0 * level.x2 - 4.0 * g * level.x4


def sturm_count(diagonal: np.ndarray, off_diagonal: float, value: float) -> int:
    """Count tridiagonal eigenvalues below ``value`` by an LDL^T recurrence."""

    tiny = np.finfo(float).tiny
    pivot = float(diagonal[0] - value)
    count = int(pivot < 0.0)
    square = off_diagonal * off_diagonal
    for entry in diagonal[1:]:
        if abs(pivot) < tiny:
            pivot = -tiny if pivot < 0.0 else tiny
        pivot = float(entry - value - square / pivot)
        count += int(pivot < 0.0)
    return count


def finite_difference_levels(
    g: float,
    count: int,
    *,
    length: float,
    interior_points: int,
    upper_hint: float,
) -> np.ndarray:
    """Compute low levels on [-L,L] with a three-point kinetic stencil."""

    require(interior_points >= 101, "too few finite-difference points")
    step = 2.0 * length / (interior_points + 1)
    x = np.linspace(-length + step, length - step, interior_points)
    off_diagonal = -1.0 / step**2
    diagonal = 2.0 / step**2 + x**2 + g * x**4

    lower = 0.0
    upper = max(2.0, float(upper_hint) + 2.0)
    while sturm_count(diagonal, off_diagonal, upper) < count:
        upper *= 2.0

    levels = np.empty(count)
    for index in range(count):
        lo = lower
        hi = upper
        for _ in range(58):
            middle = 0.5 * (lo + hi)
            if sturm_count(diagonal, off_diagonal, middle) <= index:
                lo = middle
            else:
                hi = middle
        levels[index] = 0.5 * (lo + hi)
    return levels


def romberg_finite_difference(
    g: float,
    count: int,
    *,
    length: float,
    coarse_points: int,
    upper_hint: float,
) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    """Return three meshes and an h^2/h^4-cancelled extrapolation."""

    middle_points = 2 * coarse_points + 1
    fine_points = 2 * middle_points + 1
    coarse = finite_difference_levels(
        g,
        count,
        length=length,
        interior_points=coarse_points,
        upper_hint=upper_hint,
    )
    middle = finite_difference_levels(
        g,
        count,
        length=length,
        interior_points=middle_points,
        upper_hint=upper_hint,
    )
    fine = finite_difference_levels(
        g,
        count,
        length=length,
        interior_points=fine_points,
        upper_hint=upper_hint,
    )
    second_order_coarse = (4.0 * middle - coarse) / 3.0
    second_order_fine = (4.0 * fine - middle) / 3.0
    extrapolated = (16.0 * second_order_fine - second_order_coarse) / 15.0
    return middle, fine, second_order_fine, extrapolated


def bender_wu_ground(order: int) -> list[Fraction]:
    """Return exact weak-coupling coefficients for the ground-state energy."""

    require(order >= 1, "order must be positive")
    polynomials: list[list[Fraction]] = [[Fraction(1)]]
    energies = [Fraction(1)] + [Fraction(0) for _ in range(order)]

    for k in range(1, order + 1):
        coefficients = [Fraction(0) for _ in range(2 * k + 2)]
        for m in range(2 * k, 0, -1):
            source = Fraction(0)
            for j in range(1, k):
                previous = polynomials[k - j]
                if m < len(previous):
                    source += energies[j] * previous[m]
            previous_order = polynomials[k - 1]
            if m >= 2 and m - 2 < len(previous_order):
                source -= previous_order[m - 2]
            coefficients[m] = (
                source
                + 2 * (m + 1) * (2 * m + 1) * coefficients[m + 1]
            ) / (4 * m)
        energies[k] = -2 * coefficients[1]
        polynomials.append(coefficients[:-1])
    return energies


def optimal_ground_partial_sum(g: float, coefficients: list[Fraction]) -> tuple[int, float]:
    """Truncate the raw asymptotic series at its smallest nonconstant term."""

    require(g > 0.0, "optimal truncation is requested only for g > 0")
    terms = [abs(float(coefficient) * g**k) for k, coefficient in enumerate(coefficients)]
    best_order = min(range(1, len(terms)), key=terms.__getitem__)
    value = sum(float(coefficients[k]) * g**k for k in range(best_order + 1))
    return best_order, value


def turning_point_squared(energy: float, g: float) -> float:
    """Return the positive real turning point squared without cancellation."""

    if g == 0.0:
        return energy
    return 2.0 * energy / (1.0 + math.sqrt(1.0 + 4.0 * g * energy))


def wkb_action(energy: float, g: float, nodes: int = 160) -> float:
    """Evaluate the allowed-region action by smooth Gauss--Legendre quadrature."""

    a2 = turning_point_squared(energy, g)
    abscissae, weights = np.polynomial.legendre.leggauss(nodes)
    theta = 0.25 * PI * (abscissae + 1.0)
    integrand = np.cos(theta) ** 2 * np.sqrt(
        1.0 + g * a2 * (1.0 + np.sin(theta) ** 2)
    )
    return 0.5 * PI * a2 * float(weights @ integrand)


def wkb_level(n: int, g: float) -> float:
    """Solve the leading Bohr--Sommerfeld condition for level n."""

    target = PI * (n + 0.5)
    lo = 0.0
    hi = max(2.0 * n + 1.0, 1.0)
    while wkb_action(hi, g) < target:
        hi *= 2.0
    for _ in range(60):
        middle = 0.5 * (lo + hi)
        if wkb_action(middle, g) < target:
            lo = middle
        else:
            hi = middle
    return 0.5 * (lo + hi)


def pure_quartic_data(count: int = 4) -> tuple[np.ndarray, np.ndarray]:
    """Return pure-quartic energies and <y^2> from Fourier--Galerkin blocks."""

    candidates: list[tuple[float, float]] = []
    for parity in ("even", "odd"):
        energies, vectors, _, x2, _ = fourier_block(
            parity=parity,
            quadratic=0.0,
            quartic=1.0,
            length=7.0,
            basis_size=72,
        )
        for local_index in range((count + 1) // 2 + 1):
            vector = vectors[:, local_index]
            candidates.append(
                (float(energies[local_index]), float(vector @ (x2 @ vector)))
            )
    candidates.sort()
    return (
        np.array([item[0] for item in candidates[:count]]),
        np.array([item[1] for item in candidates[:count]]),
    )


@dataclass
class CalibrationRow:
    """One coupling's cross-method output."""

    g: float
    galerkin: np.ndarray
    finite_difference: np.ndarray
    wkb_ground: float
    wkb_excited: float
    basis_shift: float
    box_shift: float
    fd_shift: float
    cross_shift: float
    virial: float


def calibrate(
    couplings: Iterable[float],
    *,
    count: int,
    basis_size: int,
    coarse_points: int,
) -> list[CalibrationRow]:
    """Run all default numerical audits."""

    rows: list[CalibrationRow] = []
    for g in couplings:
        length = box_length(g, 8.0)
        final = fourier_levels(
            g, count, basis_size=basis_size, length=length
        )
        smaller = fourier_levels(
            g, count, basis_size=basis_size - 12, length=length
        )
        shorter = fourier_levels(
            g,
            count,
            basis_size=basis_size,
            length=box_length(g, 7.0),
        )
        galerkin = np.array([level.energy for level in final])
        basis_shift = float(
            np.max(np.abs(galerkin - np.array([level.energy for level in smaller])))
        )
        box_shift = float(
            np.max(np.abs(galerkin - np.array([level.energy for level in shorter])))
        )
        _, _, second_order_fine, extrapolated = romberg_finite_difference(
            g,
            count,
            length=length,
            coarse_points=coarse_points,
            upper_hint=float(galerkin[-1]),
        )
        fd_shift = float(np.max(np.abs(extrapolated - second_order_fine)))
        cross_shift = float(np.max(np.abs(extrapolated - galerkin)))
        virial = max(abs(virial_residual(level, g)) for level in final)
        rows.append(
            CalibrationRow(
                g=g,
                galerkin=galerkin,
                finite_difference=extrapolated,
                wkb_ground=wkb_level(0, g),
                wkb_excited=wkb_level(min(3, count - 1), g),
                basis_shift=basis_shift,
                box_shift=box_shift,
                fd_shift=fd_shift,
                cross_shift=cross_shift,
                virial=virial,
            )
        )
    return rows


def print_report(
    rows: list[CalibrationRow],
    coefficients: list[Fraction],
    *,
    count: int,
) -> None:
    """Print a compact, reproducible calibration report."""

    print("Stable quartic oscillator calibration")
    print("H(g) = -d^2/dx^2 + x^2 + g x^4 on L^2(R), g >= 0")
    print(f"Python {platform.python_version()}; NumPy {np.__version__}")
    print()
    print("Ground-state Bender--Wu coefficients E_0(g) ~ sum e_k g^k")
    for k, coefficient in enumerate(coefficients[:5]):
        print(f"  e_{k} = {coefficient}")
    print()

    headings = ["g"] + [f"E_{n}" for n in range(count)]
    print("Fourier--Galerkin levels")
    print(" ".join(f"{heading:>16s}" for heading in headings))
    for row in rows:
        values = [f"{row.g:.6g}"] + [f"{value:.12f}" for value in row.galerkin]
        print(" ".join(f"{value:>16s}" for value in values))
    print()

    print("Leading WKB and empirical numerical audits")
    print(
        f"{'g':>10s} {'WKB E0':>14s} {'WKB E3*':>14s} "
        f"{'basis':>11s} {'box':>11s} {'FD':>11s} {'cross':>11s} {'virial':>11s}"
    )
    for row in rows:
        print(
            f"{row.g:10.4g} {row.wkb_ground:14.10f} {row.wkb_excited:14.10f} "
            f"{row.basis_shift:11.2e} {row.box_shift:11.2e} "
            f"{row.fd_shift:11.2e} {row.cross_shift:11.2e} {row.virial:11.2e}"
        )
    print(f"* E_{min(3, count - 1)} when fewer than four levels are requested.")
    print()

    for row in rows:
        if 0.0 < row.g <= 0.1:
            order, value = optimal_ground_partial_sum(row.g, coefficients)
            error = value - row.galerkin[0]
            cap = " (search cap)" if order == len(coefficients) - 1 else ""
            print(
                f"g={row.g:g}: smallest computed raw-series term occurs at "
                f"g^{order}{cap}; partial sum={value:.12f}; "
                f"difference {error:+.3e}"
            )
    print()

    pure_energy, pure_x2 = pure_quartic_data(count)
    print("Pure-quartic endpoint H_4=-d^2/dy^2+y^4")
    print(f"  e_0^(4) = {pure_energy[0]:.12f}")
    print(f"  <y^2>_0 = {pure_x2[0]:.12f}")
    for row in rows:
        if row.g >= 10.0:
            eta = row.g ** (-2.0 / 3.0)
            strong = row.g ** (1.0 / 3.0) * (
                pure_energy[0] + eta * pure_x2[0]
            )
            print(
                f"  g={row.g:g}: two-term strong E_0={strong:.12f}; "
                f"direct difference {strong - row.galerkin[0]:+.3e}"
            )


def write_csv(path: str, rows: list[CalibrationRow]) -> None:
    """Write figure-ready method-comparison data."""

    with open(path, "w", newline="", encoding="utf-8") as handle:
        writer = csv.writer(handle)
        writer.writerow(
            [
                "g",
                "E0",
                "E1",
                "E2",
                "E3",
                "wkb_E0",
                "wkb_E3",
                "basis_shift",
                "box_shift",
                "fd_shift",
                "cross_shift",
                "virial_residual",
            ]
        )
        for row in rows:
            padded = list(row.galerkin[:4]) + [math.nan] * max(0, 4 - row.galerkin.size)
            writer.writerow(
                [
                    row.g,
                    *padded[:4],
                    row.wkb_ground,
                    row.wkb_excited,
                    row.basis_shift,
                    row.box_shift,
                    row.fd_shift,
                    row.cross_shift,
                    row.virial,
                ]
            )


def parse_arguments() -> argparse.Namespace:
    """Parse command-line options."""

    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--high",
        action="store_true",
        help="use six levels, larger Galerkin blocks, and a finer coordinate grid",
    )
    parser.add_argument(
        "--csv",
        metavar="PATH",
        help="write the method-comparison rows to a CSV file",
    )
    return parser.parse_args()


def main() -> None:
    """Run the calibration and its consistency checks."""

    arguments = parse_arguments()
    couplings = [0.0, 0.01, 0.1, 1.0, 10.0, 100.0]
    count = 6 if arguments.high else 4
    basis_size = 76 if arguments.high else 60
    coarse_points = 1600 if arguments.high else 900
    coefficients = bender_wu_ground(20)
    expected = [
        Fraction(1),
        Fraction(3, 4),
        Fraction(-21, 16),
        Fraction(333, 64),
        Fraction(-30885, 1024),
    ]
    require(coefficients[:5] == expected, "Bender--Wu recurrence audit failed")

    rows = calibrate(
        couplings,
        count=count,
        basis_size=basis_size,
        coarse_points=coarse_points,
    )
    harmonic = rows[0].galerkin
    require(
        np.max(np.abs(harmonic - (2.0 * np.arange(count) + 1.0))) < 2.0e-10,
        "harmonic endpoint audit failed",
    )
    require(
        all(np.all(np.diff(row.galerkin) > 0.0) for row in rows),
        "level ordering audit failed",
    )
    for previous, current in zip(rows, rows[1:]):
        require(
            np.all(current.galerkin > previous.galerkin),
            "coupling monotonicity audit failed",
        )

    print_report(rows, coefficients, count=count)
    if arguments.csv:
        write_csv(arguments.csv, rows)
        print(f"\nWrote {arguments.csv}")


if __name__ == "__main__":
    main()