#!/usr/bin/env python3
# SPDX-License-Identifier: MIT
"""Deterministic AdS3--Vaidya checks for the Volume 24 quench page.

The calculation is intentionally bounded.  It integrates the outgoing-null
event horizon of a smooth planar AdS3--Vaidya shell, verifies regulated vacuum
and BTZ geodesic lengths, and checks the analytic thin-shell crossing and
causal-saturation conditions.  It does not claim to solve a time-dependent
extremal-surface boundary-value problem.
"""

from __future__ import annotations

import argparse
import json
import math
import platform
import sys
from pathlib import Path

import numpy as np
import scipy
from scipy.integrate import solve_ivp
from scipy.optimize import brentq


REFERENCE_ENVIRONMENT = {
    "python": "3.11.4",
    "numpy": "2.0.2",
    "scipy": "1.13.1",
    "floating_point": "IEEE-754 binary64",
}

DEFAULT_OUTPUT = Path(
    "public/data/string-holographic-black-holes/"
    "holographic-quench-benchmark.json"
)


def shell_mass(v: float | np.ndarray, width: float) -> float | np.ndarray:
    """Dimensionless final mass is one and the shell is centered at v=0."""

    return 0.5 * (1.0 + np.tanh(np.asarray(v) / width))


def horizon_rhs(v: float, z: np.ndarray, width: float) -> np.ndarray:
    """Outgoing radial null curve for ds^2=z^-2[-f dv^2-2 dv dz+dx^2]."""

    f_value = 1.0 - float(shell_mass(v, width)) * z[0] ** 2
    return np.array([-0.5 * f_value], dtype=float)


def integrate_event_horizon(
    width: float,
    rtol: float,
    atol: float,
    max_step: float,
    terminal_z: float = 1.0,
) -> solve_ivp:
    """Integrate backward from the final planar-BTZ horizon z_h=1."""

    solution = solve_ivp(
        lambda v, z: horizon_rhs(v, z, width),
        (4.0, -2.0),
        np.array([terminal_z], dtype=float),
        method="DOP853",
        rtol=rtol,
        atol=atol,
        max_step=max_step,
        dense_output=True,
    )
    if not solution.success or solution.sol is None:
        raise RuntimeError(f"event-horizon integration failed: {solution.message}")
    return solution


def max_integral_defect(solution: solve_ivp, width: float) -> float:
    """Trapezoidal integral defect on accepted ODE steps."""

    v = solution.t
    z = solution.y[0]
    rhs = np.array([horizon_rhs(v_i, np.array([z_i]), width)[0] for v_i, z_i in zip(v, z)])
    step = np.diff(v)
    defect = np.diff(z) - 0.5 * step * (rhs[:-1] + rhs[1:])
    return float(np.max(np.abs(defect)))


def regulated_geodesic_lengths(ell: float, cutoff: float) -> tuple[float, float, float]:
    """Vacuum and unit-temperature-scale BTZ lengths at finite z=cutoff."""

    vacuum = 2.0 * math.asinh(ell / (2.0 * cutoff)) + 2.0 * math.log(cutoff)
    thermal = 2.0 * math.asinh(math.sinh(ell / 2.0) / cutoff) + 2.0 * math.log(cutoff)
    return vacuum, thermal, thermal - vacuum


def environment_record() -> dict[str, object]:
    actual = {
        "python": platform.python_version(),
        "numpy": np.__version__,
        "scipy": scipy.__version__,
        "floating_point": "IEEE-754 binary64" if np.finfo(float).bits == 64 else "unexpected",
    }
    return {
        "reference": REFERENCE_ENVIRONMENT,
        "actual": actual,
        "matches_reference": actual == REFERENCE_ENVIRONMENT,
    }


def compute() -> dict[str, object]:
    width = 0.15
    ell = 1.2
    exact_delta = 2.0 * math.log(2.0 * math.sinh(ell / 2.0) / ell)
    refinements = [
        ("coarse", 1.0e-7, 1.0e-10, 0.04, 2.0e-2),
        ("medium", 1.0e-9, 1.0e-12, 0.02, 1.0e-2),
        ("fine", 1.0e-11, 1.0e-14, 0.01, 5.0e-3),
    ]
    reference = integrate_event_horizon(width, 5.0e-13, 5.0e-15, 0.002)
    reference_z0 = float(reference.sol(0.0)[0])

    rows: list[dict[str, object]] = []
    for label, rtol, atol, max_step, cutoff in refinements:
        solution = integrate_event_horizon(width, rtol, atol, max_step)
        vacuum, thermal, delta = regulated_geodesic_lengths(ell, cutoff)
        z0 = float(solution.sol(0.0)[0])
        rows.append(
            {
                "label": label,
                "rtol": rtol,
                "atol": atol,
                "max_step": max_step,
                "radial_cutoff": cutoff,
                "event_horizon_z_at_shell_center": z0,
                "event_horizon_change_from_reference": abs(z0 - reference_z0),
                "max_step_integral_defect": max_integral_defect(solution, width),
                "vacuum_geodesic_length_renormalized": vacuum,
                "thermal_geodesic_length_renormalized": thermal,
                "thermal_minus_vacuum": delta,
                "geodesic_delta_error": abs(delta - exact_delta),
            }
        )

    thin_widths = [0.15, 0.075, 0.0375]
    thin_shell_rows = []
    for thin_width in thin_widths:
        solution = integrate_event_horizon(thin_width, 1.0e-11, 1.0e-14, 0.005)
        thin_shell_rows.append(
            {
                "shell_width": thin_width,
                "event_horizon_z_at_shell_center": float(solution.sol(0.0)[0]),
                "difference_from_thin_shell_value_one": abs(float(solution.sol(0.0)[0]) - 1.0),
            }
        )

    source_99_time = float(brentq(lambda v: float(shell_mass(v, width)) - 0.99, -1.0, 1.0, xtol=1.0e-14))
    crossing_times = [0.2, 0.4, 0.59]
    crossings = [
        {
            "boundary_time": time,
            "shell_crossing_depth": time,
            "positive_half_crossing_x": math.sqrt((ell / 2.0) ** 2 - time**2),
        }
        for time in crossing_times
    ]
    saturation_time = ell / 2.0
    wrong_branch = integrate_event_horizon(width, 1.0e-9, 1.0e-12, 0.02, terminal_z=0.9)
    wrong_branch_error = abs(float(wrong_branch.sol(4.0)[0]) - 1.0)

    acceptance = {
        "fine_horizon_change_below_1e-10": rows[-1]["event_horizon_change_from_reference"] < 1.0e-10,
        "fine_geodesic_delta_error_below_1e-4": rows[-1]["geodesic_delta_error"] < 1.0e-4,
        "refinement_reduces_horizon_change": rows[0]["event_horizon_change_from_reference"] > rows[1]["event_horizon_change_from_reference"] > rows[2]["event_horizon_change_from_reference"],
        "refinement_reduces_cutoff_error": rows[0]["geodesic_delta_error"] > rows[1]["geodesic_delta_error"] > rows[2]["geodesic_delta_error"],
        "pre_quench_source_below_2e-6": float(shell_mass(-1.0, width)) < 2.0e-6,
        "thin_shell_crossing_saturates_causal_bound": math.isclose(saturation_time, ell / 2.0, rel_tol=0.0, abs_tol=1.0e-14),
        "wrong_terminal_branch_detected": wrong_branch_error > 1.0e-2,
    }

    return {
        "schema_version": 1,
        "model": {
            "metric": "ds^2=z^-2[-(1-m(v)z^2)dv^2-2dv dz+dx^2]",
            "shell_profile": "m(v)=(1+tanh(v/w))/2",
            "boundary_dimension": 2,
            "final_horizon_radius_z": 1.0,
            "shell_width_w": width,
            "boundary_interval_length_ell": ell,
            "initial_state": "Poincare AdS3 vacuum",
            "final_state": "unit-radius planar BTZ black brane",
            "integration_domain_v": [-2.0, 4.0],
            "subtraction": "subtract 2 log(1/z_cutoff) from each regulated geodesic length",
        },
        "environment": environment_record(),
        "solvers": {
            "ode": "scipy.integrate.solve_ivp, DOP853, backward from z(4)=1",
            "root": "scipy.optimize.brentq, xtol=1e-14",
            "reference_ode": {"rtol": 5.0e-13, "atol": 5.0e-15, "max_step": 0.002},
        },
        "refinements": rows,
        "thin_shell_limit": {
            "rows": thin_shell_rows,
            "analytic_extremal_curve": "vacuum semicircle z^2+x^2=(ell/2)^2 with v=t-z",
            "shell_crossings": crossings,
            "saturation_time": saturation_time,
        },
        "diagnostics": {
            "source_time_at_99_percent_final_energy": source_99_time,
            "source_mass_at_v_minus_one": float(shell_mass(-1.0, width)),
            "exact_late_thermal_minus_vacuum_length": exact_delta,
            "reference_event_horizon_z_at_shell_center": reference_z0,
            "wrong_terminal_branch_error": wrong_branch_error,
            "interpretive_boundary": "The event horizon is teleological; it is not a boundary equilibration time. The geodesic check fixes only static limits and the thin-shell crossing bound.",
        },
        "acceptance": acceptance,
        "passed": all(acceptance.values()),
    }


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--output", type=Path, default=DEFAULT_OUTPUT)
    parser.add_argument(
        "--allow-environment-mismatch",
        action="store_true",
        help="run for diagnostics while marking the serialized environment mismatch",
    )
    args = parser.parse_args()

    record = compute()
    environment_matches = bool(record["environment"]["matches_reference"])
    if not environment_matches and not args.allow_environment_mismatch:
        print("ERROR: runtime does not match the pinned reference environment", file=sys.stderr)
        print(json.dumps(record["environment"], indent=2, sort_keys=True), file=sys.stderr)
        return 2

    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(record, indent=2, sort_keys=True) + "\n", encoding="utf-8")

    print("label   rtol      cutoff   z_EH(v=0)       horizon change    geodesic error")
    for row in record["refinements"]:
        print(
            f"{row['label']:<7} {row['rtol']:<9.1e} {row['radial_cutoff']:<8.3g} "
            f"{row['event_horizon_z_at_shell_center']:<15.12f} "
            f"{row['event_horizon_change_from_reference']:<17.3e} "
            f"{row['geodesic_delta_error']:.3e}"
        )
    print(f"output: {args.output}")
    print(f"passed: {record['passed']}")
    return 0 if record["passed"] else 1


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