#!/usr/bin/env python3
"""Compute the pure-quartic spectrum from the folded TBA and an exact EQC.

The benchmark is

    H = -d^2/dq^2 + q^4  on L^2(R).

At fixed energy one for -hbar^2 d^2/dq^2 + q^4, put
X = hbar^(-1), theta = log(X), and E = X^(4/3).  The script solves the
two-component pure-quartic TBA, reconstructs the median real period, applies
the Ito-Marino-Shu exact quantization condition, and compares the first six
levels with an independent parity-resolved Rayleigh-Ritz calculation.

The real-line convolutions use a zero-padded Toeplitz FFT and analytic constant
left tails.  The principal value in the median period is evaluated by
singularity subtraction.  The reported refinement differences are empirical
error estimates, not interval certificates.

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

from __future__ import annotations

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

import numpy as np


PI = math.pi
SQRT2 = math.sqrt(2.0)
B0 = math.sqrt(PI) * math.gamma(0.25) / (2.0 * math.gamma(1.75))
MASS_1 = B0 / SQRT2
MASS_2 = B0

EPSILON_1_MINUS = math.log(2.0)
EPSILON_2_MINUS = math.log(3.0)
L_1_MINUS = math.log(1.5)
L_2_MINUS = math.log(4.0 / 3.0)


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

    if not condition:
        raise RuntimeError(message)


def stable_L(epsilon: np.ndarray) -> np.ndarray:
    """Evaluate log(1+exp(-epsilon)) without overflow."""

    return np.logaddexp(0.0, -epsilon)


def kappa_0(z: np.ndarray | complex | float) -> np.ndarray | complex | float:
    """Return kappa_0(z)=1/(pi cosh z)."""

    return 1.0 / (PI * np.cosh(z))


def kappa_1(z: np.ndarray | complex | float) -> np.ndarray | complex | float:
    """Return kappa_1(z)=sqrt(2) cosh(z)/(pi cosh(2z))."""

    return SQRT2 / PI * np.cosh(z) / np.cosh(2.0 * z)


def kappa_0_left_tail(a: np.ndarray | complex | float):
    """Compute integral_a^infinity kappa_0(u) du."""

    return 2.0 / PI * np.arctan(np.exp(-a))


def kappa_1_left_tail(a: np.ndarray | complex | float):
    """Compute integral_a^infinity kappa_1(u) du."""

    return 0.5 - np.arctan(SQRT2 * np.sinh(a)) / PI


@dataclass
class TbaSolution:
    """A converged discrete TBA solution and its metadata."""

    operator: "FoldedQuarticOperator"
    epsilon_1: np.ndarray
    epsilon_2: np.ndarray
    iterations: int
    residual: float

    @property
    def L_1(self) -> np.ndarray:
        return stable_L(self.epsilon_1)

    @property
    def L_2(self) -> np.ndarray:
        return stable_L(self.epsilon_2)


class FoldedQuarticOperator:
    """Uniform-grid Nyström discretization of the folded quartic TBA."""

    def __init__(self, cutoff: float, node_count: int):
        require(cutoff > 0.0, "cutoff must be positive")
        require(node_count >= 65, "at least 65 nodes are required")
        require(node_count % 2 == 1, "use an odd node count")

        self.cutoff = float(cutoff)
        self.node_count = int(node_count)
        self.theta = np.linspace(-cutoff, cutoff, node_count)
        self.step = 2.0 * cutoff / (node_count - 1)

        self.trapezoid_weights = np.full(node_count, self.step)
        self.trapezoid_weights[[0, -1]] *= 0.5

        # Composite Simpson weights are used only after the TBA solve, for the
        # smooth singularity-subtracted period integrals.
        self.simpson_weights = np.full(node_count, 2.0)
        self.simpson_weights[1:-1:2] = 4.0
        self.simpson_weights[[0, -1]] = 1.0
        self.simpson_weights *= self.step / 3.0

        lags = self.step * np.arange(node_count)
        self.kernel_ffts: list[np.ndarray] = []
        for kernel in (kappa_0, kappa_1):
            first_column = kernel(lags)
            first_row = kernel(-lags)
            embedding = np.concatenate(
                (first_column, np.array([0.0]), first_row[:0:-1])
            )
            require(
                embedding.size == 2 * node_count,
                "Toeplitz embedding has the wrong length",
            )
            self.kernel_ffts.append(np.fft.fft(embedding))

        a = self.theta + cutoff
        self.left_tail_1 = (
            kappa_0_left_tail(a) * L_1_MINUS
            + kappa_1_left_tail(a) * L_2_MINUS
        )
        self.left_tail_2 = (
            2.0 * kappa_1_left_tail(a) * L_1_MINUS
            + kappa_0_left_tail(a) * L_2_MINUS
        )

    def convolve_pair(self, values: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
        """Apply both truncated Toeplitz convolutions with one forward FFT."""

        require(values.shape == self.theta.shape, "convolution shape mismatch")
        padded = np.zeros(2 * self.node_count, dtype=np.complex128)
        padded[: self.node_count] = self.trapezoid_weights * values
        transformed = np.fft.fft(padded)
        answers = []
        for kernel_fft in self.kernel_ffts:
            answer = np.fft.ifft(kernel_fft * transformed)
            answers.append(answer[: self.node_count].real)
        return answers[0], answers[1]

    def dense_convolve(self, values: np.ndarray, kernel_index: int) -> np.ndarray:
        """Apply one truncated Nyström matrix directly for an audit."""

        kernel = (kappa_0, kappa_1)[kernel_index]
        differences = self.theta[:, None] - self.theta[None, :]
        return kernel(differences) @ (self.trapezoid_weights * values)

    def solve(
        self,
        *,
        tolerance: float = 2.0e-13,
        max_iterations: int = 500,
    ) -> TbaSolution:
        """Solve the positive-kernel two-component fixed-point problem."""

        theta = self.theta
        switch = 1.0 / (1.0 + np.exp(np.clip(theta, -700.0, 700.0)))
        epsilon_1 = MASS_1 * np.exp(theta) + EPSILON_1_MINUS * switch
        epsilon_2 = MASS_2 * np.exp(theta) + EPSILON_2_MINUS * switch

        residual = math.inf
        for iteration in range(1, max_iterations + 1):
            L_1 = stable_L(epsilon_1)
            L_2 = stable_L(epsilon_2)
            k0_L1, k1_L1 = self.convolve_pair(L_1)
            k0_L2, k1_L2 = self.convolve_pair(L_2)

            proposal_1 = (
                MASS_1 * np.exp(theta)
                + k0_L1
                + k1_L2
                + self.left_tail_1
            )
            proposal_2 = (
                MASS_2 * np.exp(theta)
                + 2.0 * k1_L1
                + k0_L2
                + self.left_tail_2
            )
            residual = float(
                max(
                    np.max(np.abs(proposal_1 - epsilon_1)),
                    np.max(np.abs(proposal_2 - epsilon_2)),
                )
            )
            if residual <= tolerance:
                return TbaSolution(
                    self, epsilon_1, epsilon_2, iteration, residual
                )
            epsilon_1, epsilon_2 = proposal_1, proposal_2

        raise RuntimeError(
            f"TBA iteration failed: T={self.cutoff}, N={self.node_count}, "
            f"residual={residual:.3e}"
        )


def cubic_value_derivative(
    grid: np.ndarray,
    values: np.ndarray,
    x: float,
    step: float,
) -> tuple[float, float]:
    """Return a local four-point cubic value and derivative."""

    index = int(math.floor((x - grid[0]) / step))
    index = max(1, min(values.size - 3, index))
    u = (x - grid[index]) / step
    y0, y1, y2, y3 = values[index - 1 : index + 3]

    w0 = -u * (u - 1.0) * (u - 2.0) / 6.0
    w1 = (u + 1.0) * (u - 1.0) * (u - 2.0) / 2.0
    w2 = -(u + 1.0) * u * (u - 2.0) / 2.0
    w3 = (u + 1.0) * u * (u - 1.0) / 6.0

    dw0 = -(3.0 * u * u - 6.0 * u + 2.0) / 6.0
    dw1 = (3.0 * u * u - 4.0 * u - 1.0) / 2.0
    dw2 = -(3.0 * u * u - 2.0 * u - 2.0) / 2.0
    dw3 = (3.0 * u * u - 1.0) / 6.0

    value = w0 * y0 + w1 * y1 + w2 * y2 + w3 * y3
    derivative = (dw0 * y0 + dw1 * y1 + dw2 * y2 + dw3 * y3) / step
    return float(value), float(derivative)


class QuarticQuantization:
    """Reconstruct the two periods and solve the pure-quartic EQC."""

    def __init__(self, solution: TbaSolution):
        self.solution = solution
        self.operator = solution.operator
        self.theta = self.operator.theta
        self.weights = self.operator.simpson_weights
        self.L_1 = solution.L_1
        self.L_2 = solution.L_2

    def value_and_derivative(self, values: np.ndarray, x: float):
        return cubic_value_derivative(
            self.theta, values, x, self.operator.step
        )

    def complex_period(self, x: float) -> float:
        """Return B=i s(Pi_gamma2)/hbar=epsilon_2 on the real ray."""

        return self.value_and_derivative(self.solution.epsilon_2, x)[0]

    def median_real_period(self, x: float) -> float:
        """Return the median-resummed real period divided by hbar."""

        differences = x - self.theta

        regular_kernel = np.sinh(differences) / np.cosh(2.0 * differences)
        integral_1 = float(np.sum(self.weights * regular_kernel * self.L_1))

        # Analytic constant left tail of sinh(u)/cosh(2u).
        a = x + self.operator.cutoff
        cosh_a = math.cosh(a)
        c = 1.0 / SQRT2
        integral_1 += L_1_MINUS / (2.0 * SQRT2) * math.log(
            (cosh_a + c) / (cosh_a - c)
        )

        L_x, L_prime_x = self.value_and_derivative(self.L_2, x)
        pv_values = np.empty_like(differences)
        small = np.abs(differences) < 1.0e-8
        pv_values[~small] = (
            self.L_2[~small] - L_x
        ) / np.sinh(differences[~small])
        pv_values[small] = -L_prime_x
        integral_2 = float(np.sum(self.weights * pv_values))

        # Analytic plateau/zero tails of the subtracted PV integral.
        left_a = x + self.operator.cutoff
        right_a = x - self.operator.cutoff
        integral_2 += (L_2_MINUS - L_x) * (
            -math.log(math.tanh(left_a / 2.0))
        )
        integral_2 += (-L_x) * math.log(abs(math.tanh(right_a / 2.0)))

        return (
            MASS_2 * math.exp(x)
            - 2.0 * SQRT2 / PI * integral_1
            - integral_2 / PI
        )

    def correction(self, level: int, x: float) -> float:
        """Return the parity-sensitive complex-cycle correction."""

        B = self.complex_period(x)
        return 2.0 * ((-1.0) ** level) * math.atan(math.exp(-B / 2.0))

    def residual(self, level: int, x: float) -> float:
        """Evaluate A-correction-2*pi*(n+1/2)."""

        return (
            self.median_real_period(x)
            - self.correction(level, x)
            - 2.0 * PI * (level + 0.5)
        )

    def solve_level(self, level: int) -> tuple[float, float]:
        """Bracket and bisect one monotone quantization crossing."""

        guess = math.log(2.0 * PI * (level + 0.5) / MASS_2)
        lower = guess - 0.5
        upper = guess + 0.5
        f_lower = self.residual(level, lower)
        f_upper = self.residual(level, upper)

        for _ in range(10):
            if f_lower * f_upper <= 0.0:
                break
            lower -= 0.25
            upper += 0.25
            f_lower = self.residual(level, lower)
            f_upper = self.residual(level, upper)
        require(f_lower * f_upper <= 0.0, f"failed to bracket level {level}")

        for _ in range(100):
            midpoint = 0.5 * (lower + upper)
            f_midpoint = self.residual(level, midpoint)
            if f_lower * f_midpoint <= 0.0:
                upper, f_upper = midpoint, f_midpoint
            else:
                lower, f_lower = midpoint, f_midpoint
            if upper - lower <= 2.0e-15:
                break

        theta_root = 0.5 * (lower + upper)
        return theta_root, upper - lower

    def levels(self, level_count: int = 6) -> list[dict[str, float]]:
        """Compute a list of quantized rapidities and energies."""

        rows: list[dict[str, float]] = []
        for level in range(level_count):
            theta_root, bracket_width = self.solve_level(level)
            X = math.exp(theta_root)
            energy = X ** (4.0 / 3.0)
            A = self.median_real_period(theta_root)
            B = self.complex_period(theta_root)
            correction = self.correction(level, theta_root)
            residual = A - correction - 2.0 * PI * (level + 0.5)
            rows.append(
                {
                    "level": float(level),
                    "theta": theta_root,
                    "X": X,
                    "energy": energy,
                    "A": A,
                    "B": B,
                    "correction": correction,
                    "residual": residual,
                    "bracket_width": bracket_width,
                }
            )
        return rows


def folded_fft_audit() -> None:
    """Compare both FFT kernels with their dense Nyström matrices."""

    operator = FoldedQuarticOperator(4.0, 65)
    values = np.exp(-operator.theta**2) * (1.0 + 0.15 * operator.theta)
    fft_values = operator.convolve_pair(values)
    for kernel_index, label in enumerate(("kappa_0", "kappa_1")):
        dense_values = operator.dense_convolve(values, kernel_index)
        error = float(np.max(np.abs(fft_values[kernel_index] - dense_values)))
        require(error < 3.0e-14, f"{label} FFT audit failed: {error:.3e}")
        print(f"{label} FFT versus dense max error: {error:.3e}")


def invariant_audit() -> None:
    """Check the mass ratio and the exact left-plateau algebra."""

    require(abs(MASS_2 / MASS_1 - SQRT2) < 3.0e-16, "mass ratio failed")
    plateau_1 = L_1_MINUS + L_2_MINUS
    plateau_2 = 2.0 * L_1_MINUS + L_2_MINUS
    require(
        abs(plateau_1 - EPSILON_1_MINUS) < 3.0e-16,
        "first plateau equation failed",
    )
    require(
        abs(plateau_2 - EPSILON_2_MINUS) < 3.0e-16,
        "second plateau equation failed",
    )
    print(f"b0=m2: {MASS_2:.16f}")
    print(f"m1:    {MASS_1:.16f}")
    print(f"m2/m1 minus sqrt(2): {MASS_2/MASS_1-SQRT2:.3e}")
    print(f"left plateau: ({plateau_1:.16f}, {plateau_2:.16f})")


def rayleigh_ritz_levels(
    level_count: int = 6,
    *,
    basis_size: int = 120,
    omega: float = 6.0,
) -> np.ndarray:
    """Diagonalize p^2+q^4 in parity blocks of an oscillator basis."""

    require(basis_size >= 2 * level_count + 8, "basis is too small")
    require(omega > 0.0, "omega must be positive")
    matrix = np.zeros((basis_size, basis_size), dtype=np.float64)

    for n in range(basis_size):
        matrix[n, n] = (
            omega * (2.0 * n + 1.0) / 2.0
            + 3.0 * (2.0 * n * n + 2.0 * n + 1.0) / (4.0 * omega**2)
        )
        if n + 2 < basis_size:
            amplitude = math.sqrt((n + 1.0) * (n + 2.0))
            value = (
                -omega / 2.0 + (2.0 * n + 3.0) / (2.0 * omega**2)
            ) * amplitude
            matrix[n, n + 2] = matrix[n + 2, n] = value
        if n + 4 < basis_size:
            value = math.sqrt(
                (n + 1.0) * (n + 2.0) * (n + 3.0) * (n + 4.0)
            ) / (4.0 * omega**2)
            matrix[n, n + 4] = matrix[n + 4, n] = value

    even = np.linalg.eigvalsh(matrix[0::2, 0::2])
    odd = np.linalg.eigvalsh(matrix[1::2, 1::2])
    interlaced = np.empty(level_count, dtype=np.float64)
    interlaced[0::2] = even[: (level_count + 1) // 2]
    interlaced[1::2] = odd[: level_count // 2]
    return interlaced


def direct_solver_audit(level_count: int = 6) -> np.ndarray:
    """Vary oscillator basis size and frequency before returning a reference."""

    reference = rayleigh_ritz_levels(level_count, basis_size=120, omega=6.0)
    variants = (
        rayleigh_ritz_levels(level_count, basis_size=100, omega=6.0),
        rayleigh_ritz_levels(level_count, basis_size=120, omega=5.0),
        rayleigh_ritz_levels(level_count, basis_size=120, omega=7.0),
    )
    spread = max(float(np.max(np.abs(values - reference))) for values in variants)
    require(spread < 2.0e-12, f"Rayleigh-Ritz refinement failed: {spread:.3e}")
    print(f"Rayleigh-Ritz basis/frequency spread: {spread:.3e}")
    return reference


def run_one(cutoff: float, node_count: int) -> tuple[TbaSolution, list[dict[str, float]]]:
    """Solve one grid and return its first six exact-WKB levels."""

    operator = FoldedQuarticOperator(cutoff, node_count)
    solution = operator.solve()
    quantization = QuarticQuantization(solution)
    levels = quantization.levels(6)
    print(
        f"T={cutoff:g}, N={node_count}, h={operator.step:.6e}, "
        f"iterations={solution.iterations}, residual={solution.residual:.3e}"
    )
    return solution, levels


def print_level_table(
    cutoff_short: list[dict[str, float]],
    medium: list[dict[str, float]],
    mesh_reference: list[dict[str, float]],
    fine: list[dict[str, float]],
    direct: np.ndarray,
) -> None:
    """Print the spectrum with separate cutoff and mesh refinements."""

    print("\nPure-quartic spectrum")
    print(
        "n  parity  X=hbar^-1         E_TBA                E_direct             "
        "|cutoff shift|  |mesh shift|  |TBA-direct|"
    )
    rows = zip(cutoff_short, medium, mesh_reference, fine, direct)
    for index, row in enumerate(rows):
        short_row, medium_row, mesh_row, fine_row, direct_energy = row
        parity = "even/N" if index % 2 == 0 else "odd/D "
        print(
            f"{index:d}  {parity}  {fine_row['X']:.13f}  "
            f"{fine_row['energy']:.15f}  {direct_energy:.15f}  "
            f"{abs(medium_row['energy']-short_row['energy']):.3e}        "
            f"{abs(fine_row['energy']-mesh_row['energy']):.3e}      "
            f"{abs(fine_row['energy']-direct_energy):.3e}"
        )

    print("\nFine-grid period audit")
    print("n  B=epsilon_2       A=median period    parity correction   EQC residual")
    for row in fine[:4]:
        print(
            f"{int(row['level']):d}  {row['B']:.15f}  {row['A']:.15f}  "
            f"{row['correction']:+.15f}  {row['residual']:+.3e}"
        )


def print_figure_data(quantization: QuarticQuantization) -> None:
    """Print a compact theta, J_even/pi, J_odd/pi table for PGFPlots."""

    print("\n# theta  J_even_over_pi  J_odd_over_pi")
    for theta in np.linspace(-0.15, 2.38, 33):
        A = quantization.median_real_period(float(theta))
        B = quantization.complex_period(float(theta))
        delta = 2.0 * math.atan(math.exp(-B / 2.0))
        print(f"{theta:.8f}  {(A-delta)/PI:.12f}  {(A+delta)/PI:.12f}")


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--high",
        action="store_true",
        help="add a T=30, N=65537 high-resolution solve",
    )
    parser.add_argument(
        "--figure-data",
        action="store_true",
        help="print sampled quantization curves from the finest solve",
    )
    args = parser.parse_args()

    print("Pure-quartic folded-TBA audit")
    print(f"Python: {platform.python_version()}")
    print(f"NumPy:  {np.__version__}")
    invariant_audit()
    folded_fft_audit()
    direct = direct_solver_audit(6)

    _, mesh_coarse = run_one(30.0, 16385)
    _, cutoff_short = run_one(22.5, 24577)
    medium_solution, medium = run_one(30.0, 32769)
    fine_solution, fine = medium_solution, medium
    mesh_reference = mesh_coarse
    if args.high:
        fine_solution, fine = run_one(30.0, 65537)
        mesh_reference = medium

    print_level_table(cutoff_short, medium, mesh_reference, fine, direct)
    if args.figure_data:
        print_figure_data(QuarticQuantization(fine_solution))
    print("\nAll pure-quartic checks passed.")


if __name__ == "__main__":
    try:
        main()
    except (RuntimeError, FloatingPointError, ValueError) as error:
        print(f"ERROR: {error}", file=sys.stderr)
        raise SystemExit(1) from error