#!/usr/bin/env python3
"""Audit a uniform-grid TBA solver and its main numerical passports.

The calibration is the reflection-symmetric real A2 minimal-chamber
equation

    epsilon(theta) = exp(theta) - K * L(theta),
    K(theta) = 1 / (2*pi*cosh(theta)),
    L(theta) = log(1 + exp(-epsilon(theta))).

The script uses an endpoint-weighted trapezoidal Nystrom rule on [-T,T].
The Toeplitz matrix-vector product is embedded in a length-2N circulant,
so the FFT computes a *linear* convolution rather than an accidental
periodic wrap.  The constant left tail is integrated analytically.  It
checks the FFT normalization against a dense matrix and prints independent
cutoff and mesh sweeps.

Requirements: CPython 3.9+ and NumPy 1.24+.

The refinement tables are convergence evidence, not interval certificates.
They do not cover complex logarithm sheets, excited roots, or kernel-pole
continuation; those require the augmented systems described on the page.
"""

from __future__ import annotations

from dataclasses import dataclass
import platform
import sys

import numpy as np


PI = np.pi
GOLDEN_RATIO = (1.0 + np.sqrt(5.0)) / 2.0
EPSILON_MINUS = -np.log(GOLDEN_RATIO)
L_MINUS = 2.0 * np.log(GOLDEN_RATIO)


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 real overflow."""

    result = np.empty_like(epsilon, dtype=np.float64)
    positive = epsilon >= 0.0
    result[positive] = np.log1p(np.exp(-epsilon[positive]))
    result[~positive] = (
        -epsilon[~positive] + np.log1p(np.exp(epsilon[~positive]))
    )
    return result


def kernel(theta: np.ndarray) -> np.ndarray:
    """Return K(theta)=1/(2*pi*cosh(theta))."""

    return 1.0 / (2.0 * PI * np.cosh(theta))


@dataclass
class GridOperator:
    """Uniform-grid quadrature and its Toeplitz FFT embedding."""

    cutoff: float
    node_count: int

    def __post_init__(self) -> None:
        require(self.cutoff > 0.0, "cutoff must be positive")
        require(self.node_count >= 5, "at least five nodes are required")
        require(self.node_count % 2 == 1, "use an odd node count so theta=0 is a node")

        self.theta = np.linspace(-self.cutoff, self.cutoff, self.node_count)
        self.step = 2.0 * self.cutoff / (self.node_count - 1)
        self.weights = np.full(self.node_count, self.step)
        self.weights[[0, -1]] *= 0.5

        lags = self.step * np.arange(self.node_count)
        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 * self.node_count,
            "Toeplitz embedding has the wrong length",
        )
        self.kernel_fft = np.fft.fft(embedding)

        # If L(theta') is replaced by its exact left plateau L_MINUS for
        # theta'<-T, this is its analytic convolution contribution.
        self.left_tail = (
            L_MINUS
            / PI
            * np.arctan(np.exp(-(self.theta + self.cutoff)))
        )

    def convolve(self, values: np.ndarray) -> np.ndarray:
        """Apply the truncated weighted convolution by a linear 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.weights * values
        transformed = np.fft.fft(padded)
        answer = np.fft.ifft(self.kernel_fft * transformed)
        return answer[: self.node_count].real

    def dense_convolve(self, values: np.ndarray) -> np.ndarray:
        """Apply the same Nystrom matrix directly for an independent audit."""

        differences = self.theta[:, None] - self.theta[None, :]
        return kernel(differences) @ (self.weights * values)

    def fixed_point_map(self, epsilon: np.ndarray) -> np.ndarray:
        """Evaluate the finite-interval map with the analytic left tail."""

        return np.exp(self.theta) - self.convolve(stable_L(epsilon)) - self.left_tail


@dataclass
class SolveResult:
    """One converged discrete solution and its audit metadata."""

    operator: GridOperator
    epsilon: np.ndarray
    iterations: int
    residual: float

    @property
    def epsilon_zero(self) -> float:
        return float(self.epsilon[self.operator.node_count // 2])


def solve(
    cutoff: float,
    node_count: int,
    *,
    damping: float = 1.0,
    tolerance: float = 2.0e-14,
    max_iterations: int = 500,
) -> SolveResult:
    """Solve the scalar calibration by damped Picard iteration."""

    require(0.0 < damping <= 1.0, "damping must lie in (0,1]")
    operator = GridOperator(cutoff, node_count)
    theta = operator.theta
    epsilon = np.exp(theta) + EPSILON_MINUS / (1.0 + np.exp(theta))

    residual = np.inf
    for iteration in range(1, max_iterations + 1):
        proposal = operator.fixed_point_map(epsilon)
        residual = float(np.max(np.abs(epsilon - proposal)))
        epsilon = (1.0 - damping) * epsilon + damping * proposal
        if residual <= tolerance:
            final_residual = float(
                np.max(np.abs(epsilon - operator.fixed_point_map(epsilon)))
            )
            if final_residual <= tolerance:
                return SolveResult(operator, epsilon, iteration, final_residual)

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


def fft_normalization_audit() -> None:
    """Compare FFT Toeplitz products with dense Nystrom matrices."""

    operator = GridOperator(4.0, 65)
    test_values = np.exp(-operator.theta**2) * (1.0 + 0.2 * operator.theta)
    fft_values = operator.convolve(test_values)
    dense_values = operator.dense_convolve(test_values)
    error = float(np.max(np.abs(fft_values - dense_values)))
    require(error < 2.0e-15, f"FFT normalization/ordering error: {error:.3e}")
    print(f"FFT versus dense max error: {error:.3e}")

    # The physical calibration kernel is even, so it cannot expose a
    # reversed lag convention. Audit an asymmetric synthetic kernel too.
    def asymmetric_kernel(theta: np.ndarray) -> np.ndarray:
        return np.exp(-0.3 * theta**2) * (1.0 + 0.1 * theta)

    lags = operator.step * np.arange(operator.node_count)
    first_column = asymmetric_kernel(lags)
    first_row = asymmetric_kernel(-lags)
    embedding = np.concatenate(
        (first_column, np.array([0.0]), first_row[:0:-1])
    )
    padded = np.zeros(2 * operator.node_count, dtype=np.complex128)
    padded[: operator.node_count] = operator.weights * test_values
    fft_asymmetric = np.fft.ifft(np.fft.fft(embedding) * np.fft.fft(padded))
    differences = operator.theta[:, None] - operator.theta[None, :]
    dense_asymmetric = asymmetric_kernel(differences) @ (
        operator.weights * test_values
    )
    asymmetric_error = float(
        np.max(
            np.abs(
                fft_asymmetric[: operator.node_count].real - dense_asymmetric
            )
        )
    )
    require(
        asymmetric_error < 2.0e-14,
        f"FFT asymmetric lag-order error: {asymmetric_error:.3e}",
    )
    print(f"FFT asymmetric lag-order max error: {asymmetric_error:.3e}")


def plateau_audit() -> None:
    """Check the exact constant solution at theta -> -infinity."""

    residual = EPSILON_MINUS + 0.5 * L_MINUS
    require(abs(residual) < 2.0e-16, "left plateau equation failed")
    print(f"left plateau epsilon_-: {EPSILON_MINUS:.16f}")
    print(f"left plateau residual:  {residual:.3e}")


def odd_node_count(cutoff: float, target_step: float) -> int:
    """Choose an odd node count close to a requested mesh spacing."""

    intervals = int(round(2.0 * cutoff / target_step))
    if intervals % 2 == 1:
        intervals += 1
    return intervals + 1


def convergence_audit() -> None:
    """Print separate domain and mesh refinements for epsilon(0)."""

    reference = solve(20.0, 4097)
    reference_value = reference.epsilon_zero
    print("\nInternal high-resolution comparison value")
    print(
        f"T=20, N=4097, epsilon(0)={reference_value:.16f}, "
        f"iterations={reference.iterations}, residual={reference.residual:.3e}"
    )

    print("\nCutoff sweep at approximately fixed h=1/32")
    print("T       N       epsilon(0)          |change from reference|    residual")
    for cutoff in (6.0, 8.0, 10.0, 12.0):
        node_count = odd_node_count(cutoff, 1.0 / 32.0)
        result = solve(cutoff, node_count)
        print(
            f"{cutoff:4.0f}  {node_count:6d}  {result.epsilon_zero: .16f}  "
            f"{abs(result.epsilon_zero-reference_value): .3e}       "
            f"{result.residual:.3e}"
        )

    print("\nMesh sweep at fixed T=12")
    print("N       h              epsilon(0)          |change from reference|    residual")
    for node_count in (193, 385, 769, 1537):
        result = solve(12.0, node_count)
        print(
            f"{node_count:4d}  {result.operator.step: .6e}  "
            f"{result.epsilon_zero: .16f}  "
            f"{abs(result.epsilon_zero-reference_value): .3e}       "
            f"{result.residual:.3e}"
        )


def main() -> None:
    """Run all deterministic audits."""

    print("TBA iteration audit")
    print(f"Python: {platform.python_version()}")
    print(f"NumPy:  {np.__version__}")
    fft_normalization_audit()
    plateau_audit()
    convergence_audit()
    print("\nAll TBA iteration checks passed.")


if __name__ == "__main__":
    try:
        main()
    except Exception as error:
        print(f"FAIL: {error}", file=sys.stderr)
        raise