#!/usr/bin/env python3
"""Reproduce the scalar Schwarzschild ell=0 fundamental mode by two methods.

Conventions: exp(-i*omega*t), z=r/(2M), Omega=2M*omega, u=r*R.
The horizon is z=1. The selected lines are future-ingoing at the horizon
and the causal outgoing Jost line continued into the upper z half-plane.

The arbitrary-precision continued fraction uses the scalar Jaffe--Leaver
recurrence and a zero terminal ratio a_(N+1)/a_N. The independent binary64
Riccati calculation derives its horizon and infinity series directly from
the radial ODE. Each method starts from the same coarse seed 0.22-0.21i;
the ODE solver never uses a continued-fraction root or a stored target.

Requirements: Python 3.9+, mpmath 1.3; the Riccati method additionally needs
NumPy and SciPy (tested with NumPy 2.0.2 and SciPy 1.13.1). --dps controls
only the continued fraction; DOP853 always uses binary64 complex arithmetic.

Examples:
  python3 schwarzschild-scalar-lab.py --verify --json scalar-results.json
  python3 schwarzschild-scalar-lab.py --profile full --verify --json full.json
  python3 schwarzschild-scalar-lab.py --method cf --depths 50,100 --verify

Verification checks depth and precision changes, ODE refinement spread,
residual divided by a locally measured derivative, and (with --method all)
agreement between independently solved roots. Default absolute tolerances
in Omega are 1e-20 for continued-fraction convergence and 1e-8 for the ODE.
These are empirical checks, not interval enclosures or proofs of uniqueness.
The root list does not provide a complete time-domain Green function.

Exit codes: 0 for passing verification or advisory exploration; 1 for failed
verification/numerical execution; 2 for invalid input. --profile full varies
all declared controls one at a time; --profile quick uses a smaller subset.

Derivations and sources:
https://heun.xyz/advanced-ode/black-holes-holography/wronskian-recurrence-qnm-conditions/
E. W. Leaver, Proc. R. Soc. A 402 (1985), doi:10.1098/rspa.1985.0119.
"""

from __future__ import annotations

import argparse
from dataclasses import asdict, dataclass, replace
from importlib import metadata
import json
import math
from pathlib import Path
import platform
import sys

import mpmath as mp


class NumericalFailure(RuntimeError):
    """A numerical calculation could not produce its declared result."""


def cf_residual(omega: mp.mpc, depth: int) -> mp.mpc:
    """Scalar ell=0 recurrence, with a_(depth+1)/a_depth=0."""
    ratio = mp.mpc(0)
    for n in range(depth, 0, -1):
        alpha = (n + 1) * (n + 1 - 2j * omega)
        beta = -(2*n*n + (2 - 8j*omega)*n - 8*omega*omega - 4j*omega + 1)
        gamma = (n - 2j * omega)**2
        denominator = beta + alpha * ratio
        if denominator == 0:
            raise NumericalFailure("Continued-fraction ratio chart is singular.")
        ratio = -gamma / denominator
    return 8*omega*omega + 4j*omega - 1 + (1 - 2j*omega)*ratio


def solve_cf(depth: int, dps: int, seed_re: str, seed_im: str) -> mp.mpc:
    with mp.workdps(dps):
        seed = mp.mpc(seed_re, seed_im)
        root = mp.findroot(
            lambda omega: cf_residual(omega, depth),
            (seed, seed + mp.mpf("0.001")),
            tol=mp.power(10, -(dps - 10)), maxsteps=50,
        )
        return +root


def horizon_coefficients(omega: complex, count: int) -> list[complex]:
    """u=(z-1)^rho H(z-1), from the radial ODE, not the Jaffe ansatz."""
    rho = -1j * omega
    coefficients = [1 + 0j]
    for k in range(1, count):
        h_term = 2*(rho+k-1)*(rho+k-2) + rho+k-1 - 4*rho*rho - 1
        i_term = (rho+k-2)*(rho+k-3) - 6*rho*rho
        get = lambda j: coefficients[j] if j >= 0 else 0j
        denominator = k*(k + 2*rho)
        if abs(denominator) < 1e-14:
            raise NumericalFailure("Resonant horizon series is outside this laboratory.")
        coefficients.append(
            -(h_term*get(k-1) + i_term*get(k-2)
              - 4*rho*rho*get(k-3) - rho*rho*get(k-4)) / denominator
        )
    return coefficients


def infinity_coefficients(omega: complex, order: int) -> list[complex]:
    """u=exp(-rho*z)*z^(-rho)*sum(c_n*z^(-n)), from the radial ODE."""
    rho = -1j * omega
    if abs(rho) < 1e-12:
        raise NumericalFailure("Omega=0 has no outgoing wave basis in this chart.")
    coefficients = [1 + 0j]
    for n in range(order):
        a_term = n*(n-1) + (2-2*rho)*n - 2*rho*rho
        b_term = -2*(n-1)*(n-2) - (2*rho+5)*(n-1) - 2*rho - 1
        c_term = (n-2)*(n-3) + (2*rho+3)*(n-2) + (rho+1)**2
        get = lambda j: coefficients[j] if j >= 0 else 0j
        coefficients.append(
            -(a_term*get(n) + b_term*get(n-1) + c_term*get(n-2))
            / (2*rho*(n+1))
        )
    return coefficients


@dataclass(frozen=True)
class ODEControls:
    cutoff: float = 40.0
    infinity_order: int = 27
    horizon_offset: float = 1e-4
    horizon_terms: int = 30
    match: float = 4.0
    angle_degrees: float = 43.5212526
    rtol: float = 3e-13
    atol: float = 3e-14


def ode_mismatch(omega: complex, controls: ODEControls) -> complex:
    """Transport independently initialized logarithmic derivatives to match."""
    from scipy.integrate import solve_ivp

    omega = complex(omega)
    rho = -1j * omega
    start = 1.0 + controls.horizon_offset
    # Use the actual binary64 coordinate in both the series and integration.
    epsilon = start - 1.0
    h = horizon_coefficients(omega, controls.horizon_terms)
    h_value = sum(value*epsilon**k for k, value in enumerate(h))
    if h_value == 0:
        raise NumericalFailure("Horizon logarithmic-derivative chart has a pole.")
    w_horizon = rho/epsilon + sum(
        k*h[k]*epsilon**(k-1) for k in range(1, len(h))
    ) / h_value

    direction = complex(math.cos(math.radians(controls.angle_degrees)),
                        math.sin(math.radians(controls.angle_degrees)))
    remote = controls.match + controls.cutoff*direction
    c = infinity_coefficients(omega, controls.infinity_order)
    g_value = sum(value*remote**(-n) for n, value in enumerate(c))
    if g_value == 0:
        raise NumericalFailure("Infinity logarithmic-derivative chart has a pole.")
    w_infinity = -rho - rho/remote - sum(
        n*c[n]*remote**(-n) for n in range(1, len(c))
    ) / (remote*g_value)

    def radial_rhs(z, w):
        return (-w*w - w/(z*(z-1)) - omega*omega*z*z/(z-1)**2
                + 1/(z*z*(z-1)))

    horizon = solve_ivp(
        radial_rhs, (start, controls.match), [w_horizon], method="DOP853",
        rtol=controls.rtol, atol=controls.atol,
    )
    infinity = solve_ivp(
        lambda distance, w: -direction*radial_rhs(remote-distance*direction, w),
        (0, controls.cutoff), [w_infinity], method="DOP853",
        rtol=controls.rtol, atol=controls.atol,
    )
    for endpoint, solution in (("horizon", horizon), ("infinity", infinity)):
        if not solution.success:
            raise NumericalFailure(f"{endpoint} integration failed: {solution.message}")
    residual = complex(horizon.y[0, -1] - infinity.y[0, -1])
    if not math.isfinite(abs(residual)):
        raise NumericalFailure("Nonfinite Riccati residual; change the coordinate chart.")
    return residual


def solve_riccati(controls: ODEControls, seed: complex) -> dict:
    """Solve two real residual components from a coarse seed, independently."""
    from scipy.optimize import root

    def residual_pair(values):
        residual = ode_mismatch(complex(float(values[0]), float(values[1])), controls)
        return [residual.real, residual.imag]

    result = root(residual_pair, [seed.real, seed.imag], tol=1e-8)
    omega = complex(float(result.x[0]), float(result.x[1]))
    if not (0.1 < omega.real < 0.4 and -0.4 < omega.imag < -0.05):
        raise NumericalFailure("ODE solve left the ell=0 fundamental-mode search region.")
    residual = abs(ode_mismatch(omega, controls))
    step = 1e-5
    derivative = (ode_mismatch(omega+step, controls)
                  - ode_mismatch(omega-step, controls)) / (2*step)
    derivative_half = (ode_mismatch(omega+step/2, controls)
                       - ode_mismatch(omega-step/2, controls)) / step
    if abs(derivative) < 1e-8:
        raise NumericalFailure("Frequency derivative is too small for a simple-root estimate.")
    estimated_shift = residual / abs(derivative)
    # A solver status alone is not evidence of convergence. Conversely, an
    # iteration can stagnate at the integration noise floor; preserve its
    # status and assess the independently recomputed residual and refinements.
    return {
        "omega_re": repr(omega.real), "omega_im": repr(omega.imag),
        "controls": asdict(controls), "solver_success": bool(result.success),
        "solver_message": str(result.message), "function_evaluations": int(result.nfev),
        "residual": residual, "derivative_modulus": abs(derivative),
        "derivative_step": step,
        "derivative_refinement_relative": abs(derivative_half-derivative)/abs(derivative),
        "local_residual_frequency_estimate": estimated_shift,
    }


def ode_cases(base: ODEControls, profile: str) -> list[tuple[str, ODEControls]]:
    cases = [
        ("baseline", base),
        ("larger infinity cutoff", replace(base, cutoff=base.cutoff+10)),
        ("higher infinity order", replace(base, infinity_order=base.infinity_order+4)),
        ("tighter ODE tolerances", replace(base, rtol=max(3e-14, base.rtol/3), atol=base.atol/3)),
        ("shifted match point", replace(base, match=base.match+1)),
    ]
    if profile == "full":
        cases.extend([
            ("smaller infinity cutoff", replace(base, cutoff=base.cutoff-10)),
            ("lower infinity order", replace(base, infinity_order=base.infinity_order-4)),
            ("smaller horizon offset", replace(base, horizon_offset=base.horizon_offset/2)),
            ("larger horizon offset", replace(base, horizon_offset=base.horizon_offset*2)),
            ("fewer horizon terms", replace(base, horizon_terms=base.horizon_terms-4)),
            ("more horizon terms", replace(base, horizon_terms=base.horizon_terms+4)),
            ("lower contour angle", replace(base, angle_degrees=base.angle_degrees-2)),
            ("higher contour angle", replace(base, angle_degrees=base.angle_degrees+2)),
            ("inner match point", replace(base, match=base.match-1)),
            ("looser ODE tolerances", replace(base, rtol=base.rtol*3, atol=base.atol*3)),
        ])
    return cases


def validate_ode_controls(c: ODEControls) -> None:
    if not all(math.isfinite(value) for value in asdict(c).values()):
        raise ValueError("ODE controls must be finite")
    if not 5 <= c.cutoff <= 300:
        raise ValueError("infinity cutoff must be between 5 and 300")
    if not 4 <= c.infinity_order <= 80 or not 4 <= c.horizon_terms <= 80:
        raise ValueError("endpoint series orders must be between 4 and 80")
    if not 1e-7 <= c.horizon_offset <= 0.1:
        raise ValueError("horizon offset must be between 1e-7 and 0.1")
    if c.match <= 1 + c.horizon_offset:
        raise ValueError("match point must be outside the horizon start")
    if not 0 < c.angle_degrees < 90:
        raise ValueError("contour angle must lie strictly between 0 and 90 degrees")
    if c.rtol < 3e-14 or c.atol <= 0:
        raise ValueError("binary64 DOP853 requires rtol >= 3e-14 and atol > 0")


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--method", choices=("cf", "riccati", "all"), default="all")
    parser.add_argument("--profile", choices=("quick", "full"), default="quick")
    parser.add_argument("--verify", action="store_true")
    parser.add_argument("--json", type=Path, help="save full inputs, refinements and checks")
    parser.add_argument("--depths", help="strictly increasing comma-separated CF depths")
    parser.add_argument("--dps", type=int, default=60, help="CF decimal precision only")
    parser.add_argument("--seed-real", default="0.22")
    parser.add_argument("--seed-imag", default="-0.21")
    parser.add_argument("--cf-tolerance", default="1e-20")
    parser.add_argument("--ode-tolerance", type=float, default=1e-8)
    parser.add_argument("--cutoff", type=float, default=40)
    parser.add_argument("--infinity-order", type=int, default=27)
    parser.add_argument("--horizon-offset", type=float, default=1e-4)
    parser.add_argument("--horizon-terms", type=int, default=30)
    parser.add_argument("--match", type=float, default=4)
    parser.add_argument("--angle", type=float, default=43.5212526)
    parser.add_argument("--rtol", type=float, default=3e-13)
    parser.add_argument("--atol", type=float, default=3e-14)
    args = parser.parse_args()
    if args.dps < 40:
        parser.error("--dps must be at least 40 (a separate 30-digit check is also run)")
    mp.mp.dps = args.dps
    try:
        seed_mp = mp.mpc(mp.mpf(args.seed_real), mp.mpf(args.seed_imag))
        cf_tolerance = mp.mpf(args.cf_tolerance)
        if not mp.isfinite(seed_mp) or not (0.15 <= seed_mp.real <= 0.3 and -0.3 <= seed_mp.imag <= -0.1):
            raise ValueError("coarse seed must have 0.15 <= Re(Omega) <= 0.3 and -0.3 <= Im(Omega) <= -0.1")
        if not mp.isfinite(cf_tolerance) or cf_tolerance <= 0:
            raise ValueError("CF tolerance must be positive and finite")
        if not math.isfinite(args.ode_tolerance) or args.ode_tolerance <= 0:
            raise ValueError("ODE tolerance must be positive and finite")
        depth_text = args.depths or ("50,100,200,400,800,1600,3200,6400" if args.profile == "full" else "400,800,1600,3200")
        depths = [int(value) for value in depth_text.split(",")]
        if len(depths) < 2 or any(value < 10 for value in depths) or any(a >= b for a, b in zip(depths, depths[1:])):
            raise ValueError("--depths needs at least two strictly increasing integers >= 10")
        base = ODEControls(args.cutoff, args.infinity_order, args.horizon_offset,
                           args.horizon_terms, args.match, args.angle, args.rtol, args.atol)
        cases = ode_cases(base, args.profile)
        if args.method in ("all", "riccati"):
            for name, controls in cases:
                try:
                    validate_ode_controls(controls)
                except ValueError as error:
                    raise ValueError(f"{name}: {error}") from error
    except (ValueError, TypeError) as error:
        parser.error(str(error))

    versions = {"Python": platform.python_version(), "mpmath": mp.__version__}
    if args.method in ("all", "riccati"):
        try:
            import scipy  # noqa: F401; fail before starting either calculation
            for package in ("numpy", "scipy"):
                versions[package] = metadata.version(package)
        except ImportError:
            parser.error("Riccati integration requires SciPy and NumPy")
    report = {
        "model": "massless scalar Schwarzschild ell=0, fundamental n=0",
        "conventions": "exp(-i*omega*t); z=r/(2M); Omega=2M*omega; u=r*R",
        "environment": versions,
        "inputs": {key: str(value) if isinstance(value, Path) else value for key, value in vars(args).items()},
        "initial_seed": {"real": args.seed_real, "imag": args.seed_imag},
        "evidence_status": "empirical convergence and independent-method comparison; no enclosure",
        "checks": [],
    }

    def check(name, value, tolerance):
        passed = bool(mp.isfinite(value) and value <= tolerance)
        report["checks"].append({"name": name, "value": str(value), "tolerance": str(tolerance), "passed": passed})
        print(f"{'PASS' if passed else 'FAIL'}: {name} = {mp.nstr(value, 8)} (limit {mp.nstr(tolerance, 6)})", flush=True)

    print(report["model"])
    print(report["conventions"])
    print("Environment:", versions)
    print("Each method starts independently from", args.seed_real, args.seed_imag + "i", flush=True)
    try:
        cf_root = None
        if args.method in ("cf", "all"):
            roots, rows = [], []
            print("\nContinued fraction: zero terminal ratio; arbitrary-precision arithmetic", flush=True)
            for depth in depths:
                omega = solve_cf(depth, args.dps, args.seed_real, args.seed_imag)
                roots.append(omega)
                row = {"depth": depth, "dps": args.dps,
                       "omega_re": mp.nstr(omega.real, args.dps), "omega_im": mp.nstr(omega.imag, args.dps),
                       "residual": mp.nstr(abs(cf_residual(omega, depth)), args.dps)}
                rows.append(row)
                print(f"N={depth:5d} Omega={mp.nstr(omega, 38)} residual={mp.nstr(abs(cf_residual(omega, depth)), 5)}", flush=True)
            cf_root = roots[-1]
            low_precision = solve_cf(depths[-1], 30, args.seed_real, args.seed_imag)
            report["cf"] = {"rows": rows, "precision_control": {
                "dps": 30, "depth": depths[-1], "omega_re": mp.nstr(low_precision.real, 30),
                "omega_im": mp.nstr(low_precision.imag, 30)},
                "M_omega_re": mp.nstr(cf_root.real/2, args.dps), "M_omega_im": mp.nstr(cf_root.imag/2, args.dps)}
            check("CF last-depth change", abs(roots[-1]-roots[-2]), cf_tolerance)
            check("CF 30-digit versus requested-precision change", abs(cf_root-low_precision), cf_tolerance)

        if args.method in ("riccati", "all"):
            print("\nIndependent Riccati roots: DOP853 binary64; all controls in JSON", flush=True)
            rows = []
            for name, controls in cases:
                row = solve_riccati(controls, complex(seed_mp))
                row["case"] = name
                rows.append(row)
                print(f"{name}: Omega={row['omega_re']}{float(row['omega_im']):+.16g}i residual={row['residual']:.3e} local_delta={row['local_residual_frequency_estimate']:.3e} solver_success={row['solver_success']}", flush=True)
            ode_root = complex(float(rows[0]["omega_re"]), float(rows[0]["omega_im"]))
            spread = max(abs(complex(float(row["omega_re"]), float(row["omega_im"]))-ode_root) for row in rows[1:])
            max_local_error = max(row["local_residual_frequency_estimate"] for row in rows)
            max_derivative_change = max(row["derivative_refinement_relative"] for row in rows)
            observed_scale = max(spread, max_local_error)
            report["riccati"] = {"rows": rows, "refinement_spread": spread,
                "local_estimate_maximum": max_local_error,
                "observed_frequency_scale": observed_scale,
                "qualification": "residual/derivative is local linearization; refinement spread is empirical and is not a rigorous bound"}
            check("ODE frequency refinement spread", spread, args.ode_tolerance)
            check("ODE local residual/derivative frequency estimate", max_local_error, args.ode_tolerance/10)
            check("ODE derivative finite-difference refinement", max_derivative_change, 1e-3)
            if cf_root is not None:
                gap = abs(cf_root-mp.mpc(ode_root.real, ode_root.imag))
                report["independent_frequency_gap"] = mp.nstr(gap, 20)
                check("independent CF versus Riccati frequency gap", gap, args.ode_tolerance)
    except (NumericalFailure, ValueError, ZeroDivisionError, OverflowError) as error:
        report["execution_error"] = str(error)
        print("Numerical execution failed:", error, file=sys.stderr, flush=True)

    passed = "execution_error" not in report and all(item["passed"] for item in report["checks"])
    report["verification_passed"] = passed
    report["exit_code"] = 1 if "execution_error" in report or (args.verify and not passed) else 0
    print("\n" + ("verification: " if args.verify else "exploratory checks: ") + ("PASS" if passed else "FAIL"))
    print("Refinements and cross-method agreement support numerical digits only; no interval certificate.")
    if not args.verify:
        print("Exploratory check failures are advisory; use --verify for enforced checks.")
    if args.json is not None:
        args.json.parent.mkdir(parents=True, exist_ok=True)
        args.json.write_text(json.dumps(report, indent=2, allow_nan=False) + "\n", encoding="utf-8")
        print("Full record:", args.json)
    return report["exit_code"]


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