#!/usr/bin/env python3
# SPDX-License-Identifier: MIT
"""Reproducible teaching benchmarks for Volume 21 computational pages.

Each case isolates one numerical question and has an analytic or independently
integrated reference.  The models are deliberately small: they test algorithms,
branch choices, constraints, and convergence, not production black-hole data.
"""

from __future__ import annotations

import argparse
import math
import platform
import sys
from collections.abc import Callable

import numpy as np
import scipy
from scipy.integrate import solve_ivp
from scipy.linalg import eigh_tridiagonal
from scipy.optimize import root_scalar


EXPECTED = ("3.9.6", "2.0.2", "1.13.1")


class BenchmarkFailure(RuntimeError):
    """Raised when a retained verification condition fails."""


def require_environment() -> None:
    actual = (platform.python_version(), np.__version__, scipy.__version__)
    if actual != EXPECTED:
        raise BenchmarkFailure(
            "reference run requires CPython {}, NumPy {}, SciPy {}; got {}, {}, {}"
            .format(*EXPECTED, *actual)
        )


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


def sci(value: float, digits: int = 7) -> str:
    return f"{float(value):.{digits}e}"


def table(headers: list[str], rows: list[list[object]]) -> None:
    strings = [[str(cell) for cell in row] for row in [headers, *rows]]
    widths = [max(len(row[index]) for row in strings) for index in range(len(headers))]
    render = lambda row: "  ".join(cell.rjust(widths[index]) for index, cell in enumerate(row))
    print(render(strings[0]))
    print("  ".join("-" * width for width in widths))
    for row in strings[1:]:
        print(render(row))


def observed_orders(errors: list[float]) -> list[str]:
    output = ["--"]
    for coarse, fine in zip(errors, errors[1:]):
        output.append(f"{math.log(coarse / fine, 2):.3f}")
    return output


def rk4(
    rhs: Callable[[float, np.ndarray], np.ndarray],
    initial: np.ndarray,
    final_time: float,
    steps: int,
    observer: Callable[[float, np.ndarray], None] | None = None,
) -> np.ndarray:
    step = final_time / steps
    time = 0.0
    state = np.array(initial, dtype=float)
    if observer is not None:
        observer(time, state)
    for _ in range(steps):
        k1 = rhs(time, state)
        k2 = rhs(time + step / 2, state + step * k1 / 2)
        k3 = rhs(time + step / 2, state + step * k2 / 2)
        k4 = rhs(time + step, state + step * k3)
        state = state + step * (k1 + 2 * k2 + 2 * k3 + k4) / 6
        time += step
        if observer is not None:
            observer(time, state)
    return state


def run_self_consistent() -> None:
    """Select the weak-curvature root of x = Lambda/3 + epsilon x^2."""
    cosmological_term = 0.12
    epsilon = 0.35
    discriminant = 1 - 4 * epsilon * cosmological_term
    exact = (1 - math.sqrt(discriminant)) / (2 * epsilon)
    high_root = (1 + math.sqrt(discriminant)) / (2 * epsilon)
    tolerances = [1e-3, 1e-6, 1e-9, 1e-12]
    results = []
    for tolerance in tolerances:
        solution = root_scalar(
            lambda x: x - cosmological_term - epsilon * x * x,
            bracket=(0.0, 1.0 / (2 * epsilon)),
            method="brentq",
            xtol=tolerance,
            rtol=max(4 * np.finfo(float).eps, tolerance),
        )
        require(solution.converged, "weak-branch root solve did not converge")
        root = solution.root
        residual = abs(root - cosmological_term - epsilon * root * root)
        results.append((tolerance, root, abs(root - exact), residual, solution.iterations))

    print("CASE self-consistent-static-cosmological-solutions")
    print("Equation: x = Lambda/3 + epsilon x^2, with x=H^2.")
    print(f"Inputs: Lambda/3={cosmological_term}, epsilon={epsilon}; weak branch selected by x -> Lambda/3 as epsilon -> 0.")
    table(
        ["root tolerance", "weak x", "root error", "equation residual", "iterations"],
        [[sci(t), sci(x), sci(error), sci(residual), count] for t, x, error, residual, count in results],
    )
    require(results[-1][2] < results[0][2], "weak-root error did not improve")
    require(results[-1][3] < 2e-12, "weak-root residual exceeded tolerance")
    require(high_root > 2.0, "the separated high-curvature branch was not found")
    print(f"Independent check: quadratic formula gives weak x={exact:.12f} and high x={high_root:.12f}.")
    print("Interpretation: this algebraic de Sitter test checks branch selection and residuals; it is not a construction of a four-dimensional Hadamard state.")


def run_linear_stability() -> None:
    """Evolve the low-frequency root of a fourth-order response polynomial."""
    epsilon = 0.03
    wave_number = 1.2
    discriminant = 1 - 4 * epsilon * wave_number**2
    low_z = (1 - math.sqrt(discriminant)) / (2 * epsilon)
    high_z = (1 + math.sqrt(discriminant)) / (2 * epsilon)
    omega = math.sqrt(low_z)
    final_time = 4.3
    exact = np.array([math.cos(omega * final_time), -omega * math.sin(omega * final_time)])

    def rhs(_time: float, state: np.ndarray) -> np.ndarray:
        return np.array([state[1], -low_z * state[0]])

    steps = [20, 40, 80, 160]
    rows = []
    errors = []
    for count in steps:
        maximum_invariant_drift = 0.0

        def observe(_time: float, state: np.ndarray) -> None:
            nonlocal maximum_invariant_drift
            invariant = state[0] ** 2 + state[1] ** 2 / low_z
            maximum_invariant_drift = max(maximum_invariant_drift, abs(invariant - 1))

        numerical = rk4(rhs, np.array([1.0, 0.0]), final_time, count, observe)
        error = float(np.linalg.norm(numerical - exact))
        errors.append(error)
        rows.append([count, sci(final_time / count), sci(error), sci(maximum_invariant_drift)])

    orders = observed_orders(errors)
    for row, order in zip(rows, orders):
        row.append(order)
    print("CASE linear-stability-semiclassical-solutions")
    print("Dispersion model: epsilon z^2 - z + k^2 = 0, z=omega^2; the low root is retained below the cutoff.")
    print(f"Inputs: epsilon={epsilon}, k={wave_number}, t_f={final_time}; low omega={omega:.10f}, high omega={math.sqrt(high_z):.10f}.")
    table(["steps", "Delta t", "state error", "max invariant drift", "order"], rows)
    require(all(fine < coarse for coarse, fine in zip(errors, errors[1:])), "RK4 state error did not decrease")
    require(all(3.8 < float(order) < 4.2 for order in orders[2:]), "low-mode evolution was not fourth order")
    require(errors[-1] < 2e-7, "refined low-mode error is too large")
    polynomial_residual = abs(epsilon * low_z**2 - low_z + wave_number**2)
    require(polynomial_residual < 1e-13, "response root failed its dispersion relation")
    print(f"Independent check: exact cosine evolution and polynomial residual {polynomial_residual:.3e}; high root is kept only as a cutoff diagnostic.")


def run_negative_mode() -> None:
    """Converge the unique bound state of a Pöschl-Teller test operator."""
    half_width = 8.0
    resolutions = [64, 128, 256, 512]
    exact = -1.0
    rows = []
    errors = []
    residuals = []
    for points in resolutions:
        spacing = 2 * half_width / (points + 1)
        x = np.linspace(-half_width + spacing, half_width - spacing, points)
        potential = -2.0 / np.cosh(x) ** 2
        diagonal = 2.0 / spacing**2 + potential
        off_diagonal = np.full(points - 1, -1.0 / spacing**2)
        values, vectors = eigh_tridiagonal(
            diagonal,
            off_diagonal,
            select="i",
            select_range=(0, 1),
            check_finite=True,
        )
        eigenvalue = float(values[0])
        eigenvector = vectors[:, 0]
        applied = diagonal * eigenvector
        applied[:-1] += off_diagonal * eigenvector[1:]
        applied[1:] += off_diagonal * eigenvector[:-1]
        residual = float(np.linalg.norm(applied - eigenvalue * eigenvector))
        error = abs(eigenvalue - exact)
        errors.append(error)
        residuals.append(residual)
        rows.append([points, sci(spacing), sci(eigenvalue), sci(error), sci(residual)])

    print("CASE negative-modes-thermodynamic-stability")
    print("Operator benchmark: L=-d^2/dx^2-2 sech^2(x) on [-8,8] with Dirichlet endpoints.")
    print("The infinite-line ground-state eigenvalue is -1 and there is exactly one negative eigenvalue.")
    table(["N", "Delta x", "lowest eigenvalue", "eigenvalue error", "matrix residual"], rows)
    require(all(fine < coarse for coarse, fine in zip(errors, errors[1:])), "eigenvalue error did not decrease")
    require(errors[-1] < 1.0e-4, "refined ground-state error is too large")
    require(residuals[-1] < 1.0e-11, "refined matrix residual is too large")
    require(values[1] > -2e-3, "test operator developed an unexpected second negative mode")
    print("Independent check: the analytic sech(x) eigenfunction has eigenvalue -1; this test validates the radial eigensolver, not the Schwarzschild tensor potential.")


def evaporation_rhs(_time: float, state: np.ndarray) -> np.ndarray:
    mass, angular_momentum, charge = state
    spin = angular_momentum / mass**2
    charge_ratio = charge / mass
    mass_flux = 0.04 * (1 + 0.60 * spin**2 + 0.20 * charge_ratio**2) / mass**2
    angular_flux = 0.12 * angular_momentum * (1 + 0.40 * spin**2) / mass**3
    charge_flux = 0.18 * charge * (1 + 0.30 * charge_ratio**2) / mass**2
    return -np.array([mass_flux, angular_flux, charge_flux])


def run_evaporation() -> None:
    initial = np.array([1.0, 0.45, 0.25])
    final_time = 0.6
    reference_solution = solve_ivp(
        evaporation_rhs,
        (0.0, final_time),
        initial,
        method="DOP853",
        rtol=1e-13,
        atol=1e-15,
    )
    require(reference_solution.success, "DOP853 reference evolution failed")
    reference = reference_solution.y[:, -1]
    steps = [8, 16, 32, 64]
    rows = []
    errors = []
    for count in steps:
        trajectory: list[np.ndarray] = []
        rk4(evaporation_rhs, initial, final_time, count, lambda _t, y: trajectory.append(y.copy()))
        final = trajectory[-1]
        error = float(np.linalg.norm(final - reference))
        errors.append(error)
        times = np.linspace(0.0, final_time, count + 1)
        mass_fluxes = np.array([-evaporation_rhs(t, state)[0] for t, state in zip(times, trajectory)])
        balance = abs(initial[0] - final[0] - float(np.trapezoid(mass_fluxes, times)))
        extremality = max((state[1] / state[0] ** 2) ** 2 + (state[2] / state[0]) ** 2 for state in trajectory)
        rows.append([count, sci(final_time / count), sci(final[0]), sci(final[1]), sci(final[2]), sci(error), sci(balance), f"{extremality:.7f}"])

    print("CASE rotating-charged-evaporation")
    print("Dimensionless pedagogical flux model with M(0)=1, J(0)=0.45, Q(0)=0.25; coefficients are declared proxies, not fitted greybody data.")
    table(["steps", "Delta t", "M_f", "J_f", "Q_f", "state error", "mass mismatch", "max chi^2+q^2"], rows)
    require(all(fine < coarse for coarse, fine in zip(errors, errors[1:])), "coupled evaporation error did not decrease")
    require(all(3.7 < float(order) < 4.3 for order in observed_orders(errors)[2:]), "coupled evaporation was not fourth order")
    require(errors[-1] < 1e-9, "refined evaporation state error is too large")
    require(max(float(row[-1]) for row in rows) < 1.0, "trajectory crossed the extremality bound")

    neutral_reference = (initial[0] ** 3 - 3 * 0.04 * final_time) ** (1 / 3)
    neutral = rk4(
        lambda _t, y: np.array([-0.04 / y[0] ** 2]),
        np.array([initial[0]]),
        final_time,
        256,
    )[0]
    require(abs(neutral - neutral_reference) < 2e-11, "neutral analytic limit failed")
    print(f"Independent check: setting J=Q=0 recovers M^3=M_0^3-3 alpha t with final error {abs(neutral-neutral_reference):.3e}.")


def run_two_dimensional() -> None:
    initial_mass = 2.0
    asymptotic_flux = 0.08
    switch_time = 0.7
    final_time = 3.0

    def rhs(_time: float, state: np.ndarray) -> np.ndarray:
        mass, flux, radiated = state
        del mass, radiated
        return np.array([-flux, (asymptotic_flux - flux) / switch_time, flux])

    exact_flux = asymptotic_flux * (1 - math.exp(-final_time / switch_time))
    exact_radiated = asymptotic_flux * (final_time - switch_time * (1 - math.exp(-final_time / switch_time)))
    exact = np.array([initial_mass - exact_radiated, exact_flux, exact_radiated])
    steps = [6, 12, 24, 48]
    rows = []
    errors = []
    for count in steps:
        final = rk4(rhs, np.array([initial_mass, 0.0, 0.0]), final_time, count)
        error = float(np.linalg.norm(final - exact))
        balance = abs(final[0] + final[2] - initial_mass)
        errors.append(error)
        rows.append([count, sci(final_time / count), sci(final[0]), sci(final[1]), sci(error), sci(balance)])

    print("CASE two-dimensional-backreaction-models")
    print("Flux-balance benchmark: dM/du=-F, tau dF/du=F_infinity-F, dE_rad/du=F.")
    print(f"Inputs: M0={initial_mass}, F_infinity={asymptotic_flux}, tau={switch_time}, u_f={final_time}.")
    table(["steps", "Delta u", "Bondi M_f", "flux F_f", "state error", "balance residual"], rows)
    require(all(fine < coarse for coarse, fine in zip(errors, errors[1:])), "two-dimensional state error did not decrease")
    require(all(3.7 < float(order) < 4.3 for order in observed_orders(errors)[2:]), "flux-switch evolution was not fourth order")
    require(errors[-1] < 2e-8, "refined flux-switch error is too large")
    require(max(float(row[-1]) for row in rows) < 2e-14, "Bondi balance was not preserved")
    print("Independent check: the exponential switch-on and integrated Bondi mass are analytic. The benchmark tests flux balance, not the full CGHS or RST constraints.")


def run_trace_anomaly() -> None:
    hubble = 1.0
    effective_mass_squared = -0.4
    discriminant = 9 * hubble**2 - 4 * effective_mass_squared
    growing_exponent = (-3 * hubble + math.sqrt(discriminant)) / 2
    decaying_exponent = (-3 * hubble - math.sqrt(discriminant)) / 2
    final_time = 5.0

    def rhs(_time: float, state: np.ndarray) -> np.ndarray:
        return np.array([state[1], -3 * hubble * state[1] - effective_mass_squared * state[0]])

    initial = np.array([1.0, growing_exponent])
    exact = math.exp(growing_exponent * final_time) * initial
    steps = [10, 20, 40, 80]
    rows = []
    errors = []
    measured_exponents = []
    for count in steps:
        final = rk4(rhs, initial, final_time, count)
        error = float(np.linalg.norm(final - exact))
        measured = math.log(abs(final[0] / initial[0])) / final_time
        characteristic_residual = abs(measured**2 + 3 * hubble * measured + effective_mass_squared)
        errors.append(error)
        measured_exponents.append(measured)
        rows.append([count, sci(final_time / count), sci(final[0]), sci(measured), sci(error), sci(characteristic_residual)])

    print("CASE trace-anomaly-driven-cosmology")
    print("Linearized branch test: deltaH''+3 H0 deltaH'+mu_eff^2 deltaH=0.")
    print(f"Inputs: H0={hubble}, mu_eff^2={effective_mass_squared}, t_f={final_time}; growing-mode initial data.")
    table(["steps", "Delta t", "deltaH_f", "measured exponent", "state error", "characteristic residual"], rows)
    require(all(fine < coarse for coarse, fine in zip(errors, errors[1:])), "anomaly-mode state error did not decrease")
    require(all(3.7 < float(order) < 4.3 for order in observed_orders(errors)[2:]), "anomaly-mode evolution was not fourth order")
    require(errors[-1] < 2e-7, "refined anomaly-mode error is too large")
    require(abs(measured_exponents[-1] - growing_exponent) < 5e-11, "refined growth exponent is inaccurate")
    print(f"Independent check: characteristic roots are s_+={growing_exponent:.12f}, s_-={decaying_exponent:.12f}; the negative mu_eff^2 branch grows.")
    print("Changing the anomaly-dependent effective mass to a positive value makes both root real parts negative; the full nonlinear anomaly equation has additional validity conditions.")


CASES = {
    "self-consistent-static-cosmological-solutions": run_self_consistent,
    "linear-stability-semiclassical-solutions": run_linear_stability,
    "negative-modes-thermodynamic-stability": run_negative_mode,
    "rotating-charged-evaporation": run_evaporation,
    "two-dimensional-backreaction-models": run_two_dimensional,
    "trace-anomaly-driven-cosmology": run_trace_anomaly,
}


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--case", choices=[*CASES, "all"], default="all")
    args = parser.parse_args()
    require_environment()
    selected = CASES.items() if args.case == "all" else [(args.case, CASES[args.case])]
    for index, (_name, runner) in enumerate(selected):
        if index:
            print()
        runner()


if __name__ == "__main__":
    try:
        main()
    except BenchmarkFailure as error:
        print(f"BENCHMARK FAILED: {error}", file=sys.stderr)
        raise SystemExit(1) from error