#!/usr/bin/env python3
"""Pinned coupled Kerr check; prints deterministic JSON, does not fetch mode data.

Units: w=GM*omega, a_star=a/(GM), time dependence exp(-i*omega*t).
qnm 0.4.4 uses the same dimensionless w and a_star. The independent
radial coefficients below translate Leaver1985 Eqs25-26 from W=2*w,
a_L=a_star/2. No qnm coefficient or CF function is used by that check.
"""
import os
from pathlib import Path

HERE = Path(__file__).resolve().parent
os.environ["QNMCACHEDIR"] = str(HERE / "unused-qnm-cache")
os.environ["NUMBA_DISABLE_JIT"] = "1"
for name in ("OPENBLAS_NUM_THREADS", "OMP_NUM_THREADS", "MKL_NUM_THREADS",
             "VECLIB_MAXIMUM_THREADS"):
    os.environ[name] = "1"

import hashlib
import importlib.metadata
import json
import platform
import sys
import numpy as np
import scipy
from mpmath import mp
import qnm
from qnm.nearby import NearbyRootFinder

S, ELL, M, N_INV = -2, 2, 1, 0
mp.dps = 60


def require_scope(s, ell, m, overtone, a_star):
    """This file is one finite benchmark, not a general Kerr CLI."""
    if (s, ell, m, overtone, a_star) != (-2, 2, 1, 0, .4):
        raise ValueError("Only s=-2, ell=2, m=1, n=0, a_star=0.4 is supported")


def accept_angular(previous_overlap, extended_relative_residual):
    return previous_overlap > .99 and extended_relative_residual < 1e-11


def pair(z):
    z = complex(z)
    return [float(z.real), float(z.imag)]


def mp_pair(z):
    return [mp.nstr(mp.re(z), 35), mp.nstr(mp.im(z), 35)]


def mpc(z):
    return mp.mpc(repr(float(np.real(z))), repr(float(np.imag(z))))


def leaver_coeffs(w, a_star, A):
    """Original 1985 Eq26, not the package's D_coeffs implementation."""
    W, a = 2 * w, a_star / 2
    b = mp.sqrt(1 - 4*a*a)
    q = W/2 - a*M
    return [1-S-1j*W-2j*q/b,
            -4+2*(2+b)*1j*W+4j*q/b,
            S+3-3j*W-2j*q/b,
            W*W*(4+2*b-a*a)-2*a*M*W-S-1-A+(2+b)*1j*W
            +(4*W+2j)*q/b,
            S+1-2*W*W-(2*S+3)*1j*W-(4*W+2j)*q/b]


def radial_fraction(w, a_star, A, depth):
    """F0 with zero tail T_{depth+1}=-a_{depth+1}/a_depth=0."""
    c = leaver_coeffs(w, a_star, A)
    tail = mp.mpc(0)
    for j in range(depth, 0, -1):
        al = j*j+(c[0]+1)*j+c[0]
        be = -2*j*j+(c[1]+2)*j+c[3]
        ga = j*j+(c[2]-3)*j+c[4]-c[2]+2
        tail = ga/(be-al*tail)
    return c[3]-c[0]*tail


def angular_fraction(c, A, depth):
    """Leaver1985 printed290 Eqs18-21, independent of angular matrices.

    S=e^(cu)(1+u)^k1(1-u)^k2 sum a_j(1+u)^j. For this branch
    ell=k1+k2=2 at c=0, so inversion0 starts at A=4.
    """
    k1, k2 = mp.mpf(abs(M-S))/2, mp.mpf(abs(M+S))/2
    def coefficients(j):
        al = -2*(j+1)*(j+2*k1+1)
        be = (j*(j-1)+2*j*(k1+k2+1-2*c)
              -2*c*(2*k1+S+1)+(k1+k2)*(k1+k2+1)
              -c*c-S*(S+1)-A)
        ga = 2*c*(j+k1+k2+S)
        return al, be, ga
    tail = mp.mpc(0)
    for j in range(depth, 0, -1):
        al, be, ga = coefficients(j)
        tail = ga/(be-al*tail)
    al, be, _ = coefficients(0)
    return be-al*tail


def angular_matrix(c, l_max):
    """Independent cos(theta) recurrence, with untruncated intermediate paths.

    H=diag(l(l+1)-s(s+1))+2cs Z-c^2 Z^2, projected only after
    squaring Z on l_min..l_max+1. The extra state matters at the last
    diagonal element: squaring an already truncated Z is incorrect.
    """
    ls = np.arange(max(abs(S), abs(M)), l_max+2)
    Z = np.zeros((len(ls), len(ls)))
    for j, l in enumerate(ls):
        Z[j, j] = -S*M/(l*(l+1))
        if j+1 < len(ls):
            u = l+1
            z = np.sqrt((u*u-M*M)*(u*u-S*S)/(u*u*(4*u*u-1)))
            Z[j, j+1] = Z[j+1, j] = z
    H = np.diag(ls*(ls+1)-S*(S+1))+2*c*S*Z-c*c*(Z@Z)
    return H[:-1, :-1]


def normed(v):
    return v/np.linalg.norm(v)


def angular_diagnostics(w, a, A, C, l_max, previous_C=None):
    C = normed(C)
    H = angular_matrix(a*w, l_max)
    package_H = qnm.angular.M_matrix(S, a*w, M, l_max)
    vals, vecs = np.linalg.eig(H)
    overlaps = np.abs(vecs.conj().T@C)
    pick = int(np.argmax(overlaps))
    v = normed(vecs[:, pick])
    other = np.delete(vals, pick)
    finite_r = H@C-A*C
    extended_l_max = l_max+4
    H_ext = angular_matrix(a*w, extended_l_max)
    padded = np.pad(C, (0, 4))
    extended_r = H_ext@padded-A*padded
    # Normalize by the two equation terms, not ||H_ext|| which grows
    # with the arbitrary spectral cutoff and could hide leakage.
    scale = np.linalg.norm(H_ext@padded)+abs(A)*np.linalg.norm(padded)
    d = {
        "ell_max": l_max,
        "matrix_formula_max_abs_difference": float(np.max(np.abs(H-package_H))),
        "independent_eigenvalue_difference": float(abs(vals[pick]-A)),
        "independent_eigenvector_overlap": float(min(1, abs(np.vdot(v, C)))),
        "nearest_other_eigenvalue_gap": float(np.min(abs(other-vals[pick]))),
        "coefficient_norm": float(np.linalg.norm(C)),
        "last_two_coefficients_l2": float(np.linalg.norm(C[-2:])),
        "finite_residual_l2": float(np.linalg.norm(finite_r)),
        "extended_ell_max": extended_l_max,
        "extended_residual_l2": float(np.linalg.norm(extended_r)),
        "extended_relative_residual": float(np.linalg.norm(extended_r)/scale),
        "omitted_rows_residual_l2": float(np.linalg.norm(extended_r[len(C):])),
    }
    if previous_C is not None:
        previous_C = normed(previous_C)
        prev_overlaps = np.abs(vecs.conj().T@previous_C)
        prev_pick = int(np.argmax(prev_overlaps))
        d.update({
            "previous_step_overlap": float(min(1, abs(np.vdot(previous_C, C)))),
            "largest_previous_overlap_selects_same_branch": bool(prev_pick == pick),
        })
    return d


def solve(a, l_max, guess, A_guess, nr_min=200):
    finder = NearbyRootFinder(a=float(a), s=S, m=M, n_inv=N_INV,
                             l_max=l_max, omega_guess=guess,
                             A_closest_to=A_guess, tol=1e-11,
                             cf_tol=1e-13, Nr_min=nr_min, Nr_max=8000)
    w = finder.do_solve()
    if w is None or not finder.solved:
        raise RuntimeError(str(finder.opt_res))
    A, C = finder.A, normed(finder.C)
    F, err, used = qnm.radial.leaver_cf_inv_lentz(
        w, float(a), S, M, A, N_INV, tol=1e-13, N_min=nr_min, N_max=8000)
    if abs(F) > 5e-11 or used >= 7999:
        raise RuntimeError("Unacceptable freshly evaluated radial residual")
    record = {"a_star": float(a), "w": pair(w), "A": pair(A),
              "ell_max": l_max, "fresh_F0": pair(F),
              "fresh_abs_F0": float(abs(F)),
              "lentz_update_error_indicator": float(err),
              "lentz_iterations": int(used),
              "root_function_evaluations": int(finder.opt_res.nfev)}
    return w, A, C, record


def main():
    if len(sys.argv) != 1:
        raise SystemExit("This finite no-input benchmark accepts no CLI arguments")
    require_scope(S, ELL, M, N_INV, .4)
    w, A = .37367-.08896j, 4+0j
    prev_C = None
    path = []
    # Explicit root solves, not KerrSpinSeq interpolation or modes_cache.
    for a in np.linspace(0, .4, 9):
        w, A, C, record = solve(a, 12, w, A)
        record["angular"] = angular_diagnostics(w, a, A, C, 12, prev_C)
        if prev_C is not None:
            assert record["angular"]["largest_previous_overlap_selects_same_branch"]
            assert record["angular"]["previous_step_overlap"] > .99
        path.append(record)
        prev_C = C

    resolution = []
    # Re-solve the coupled problem at each angular cutoff, keeping radial
    # settings fixed; then evaluate vector leakage in an enlarged operator.
    for l_max in (4, 6, 8, 10, 12, 16, 20):
        wl, Al, Cl, record = solve(.4, l_max, w, A)
        record["angular"] = angular_diagnostics(wl, .4, Al, Cl, l_max)
        record["w_difference_from_ell_max_12"] = float(abs(wl-w))
        record["A_difference_from_ell_max_12"] = float(abs(Al-A))
        resolution.append(record)

    # Evaluate the same accepted w,A with higher minimum Lentz depths and
    # fresh independent arbitrary-precision backward fractions. These are
    # residual checks of double-precision input, NOT 60-digit roots.
    radial = []
    last = None
    for depth in (128, 256, 512, 1024):
        F_mp = radial_fraction(mpc(w), mp.mpf(".4"), mpc(A), depth)
        F_q, err, used = qnm.radial.leaver_cf_inv_lentz(
            w, .4, S, M, A, 0, tol=1e-13, N_min=depth, N_max=8000)
        record = {"depth": depth, "mp_F0": mp_pair(F_mp),
                  "mp_abs_F0": mp.nstr(abs(F_mp), 20),
                  "package_F0": pair(F_q), "package_abs_F0": float(abs(F_q)),
                  "package_lentz_iterations": int(used),
                  "package_lentz_update_error_indicator": float(err),
                  "mp_change_from_previous_depth": None if last is None
                  else mp.nstr(abs(F_mp-last), 20)}
        radial.append(record)
        last = F_mp

    angular_series = []
    c_mp = mp.mpf(".4")*mpc(w)
    for depth in (32, 64, 128):
        A_cf = mp.findroot(lambda z: angular_fraction(c_mp, z, depth),
                           (mpc(A)-mp.mpf(".001"), mpc(A)+mp.mpf(".001")),
                           tol=mp.mpf("1e-50"))
        angular_series.append({
            "depth": depth, "fixed_w": pair(w), "A_cf": mp_pair(A_cf),
            "abs_difference_from_package_A": mp.nstr(abs(A_cf-mpc(A)), 20),
            "own_fraction_abs_residual": mp.nstr(abs(angular_fraction(c_mp, A_cf, depth)), 20),
            "depth_256_abs_residual": mp.nstr(abs(angular_fraction(c_mp, A_cf, 256)), 20),
        })
    assert abs(A_cf-mpc(A)) < mp.mpf("1e-11")

    # Negative control 1: a different exact finite-block eigenpair can have
    # a tiny equation residual but is not the tracked ell=2 branch.
    H = angular_matrix(.4*w, 12)
    vals, vecs = np.linalg.eig(H)
    correct = int(np.argmin(abs(vals-A)))
    wrong = min((j for j in range(len(vals)) if j != correct),
                key=lambda j: abs(vals[j]-A))
    wrong_A, wrong_C = vals[wrong], normed(vecs[:, wrong])
    # The last previous-step overlap recorded on the path uses the a=.35
    # vector; reconstruct that vector by a direct angular solve, not a cache.
    before = path[-2]
    before_w = complex(*before["w"])
    before_A = complex(*before["A"])
    _, before_C = qnm.angular.C_and_sep_const_closest(before_A, S,
                         before["a_star"]*before_w, M, 12)
    wrong_d = angular_diagnostics(w, .4, wrong_A, wrong_C, 12, before_C)
    wrong_rejected = not accept_angular(wrong_d["previous_step_overlap"],
                                      wrong_d["extended_relative_residual"])
    # Negative control 2: the finite residual alone would pass a grossly
    # under-resolved ell_max=4 vector. The extended residual must reject it.
    coarse = resolution[0]["angular"]
    coarse_rejected = not accept_angular(1., coarse["extended_relative_residual"])
    assert wrong_d["finite_residual_l2"] < 1e-11 and wrong_rejected
    assert coarse["finite_residual_l2"] < 1e-11 and coarse_rejected
    assert accept_angular(path[-1]["angular"]["previous_step_overlap"],
                          path[-1]["angular"]["extended_relative_residual"])
    unsupported = []
    for name, request in [("changed_spin", (-1, 2, 1, 0, .4)),
                          ("near_extremal", (-2, 2, 1, 0, .999999))]:
        rejected = False
        try:
            require_scope(*request)
        except ValueError:
            rejected = True
        assert rejected
        unsupported.append({"control": name, "request": list(request),
                            "rejected": rejected})
    negative_controls = {
        "wrong_angular_branch": {"A": pair(wrong_A), "diagnostics": wrong_d,
                                  "rejected": wrong_rejected},
        "false_convergence_from_tiny_finite_residual": {
            "ell_max": 4, "finite_residual_l2": coarse["finite_residual_l2"],
            "extended_relative_residual": coarse["extended_relative_residual"],
            "rejected": coarse_rejected},
        "unsupported_scope": unsupported,
        "acceptance_rule": "previous-step overlap >0.99 AND extended relative residual <1e-11",
        "accepted_tracked_final_branch": True,
    }

    mp_D = leaver_coeffs(mpc(w), mp.mpf(".4"), mpc(A))
    q_D = qnm.radial.D_coeffs(w, .4, S, M, A)
    D_difference = max(abs(complex(x)-y) for x, y in zip(mp_D, q_D))
    W_L, a_L = 2*w, .2
    printed_w = .815958-.174514j
    printed_A = 3.76757+.05324j
    reference = {"source": "Leaver1985 Table3 printed296, fundamental ell2 m1",
                 "a_L": a_L, "w_L": pair(W_L),
                 "table_w_L": pair(printed_w), "table_A": pair(printed_A),
                 "w_L_minus_table": pair(W_L-printed_w),
                 "A_minus_table": pair(A-printed_A),
                 "w_L_rounds_to_table": bool(round(W_L.real, 6)==printed_w.real
                                              and round(W_L.imag, 6)==printed_w.imag),
                 "A_rounds_to_table": bool(round(A.real, 5)==printed_A.real
                                            and round(A.imag, 5)==printed_A.imag)}
    assert reference["w_L_rounds_to_table"] and reference["A_rounds_to_table"]
    assert D_difference < 1e-12
    assert abs(last) < mp.mpf("5e-11")
    assert resolution[-1]["angular"]["extended_relative_residual"] < 1e-11
    modules = (qnm.radial, qnm.angular, qnm.nearby)
    out = {
        "schema": "frobenius-kerr-v1",
        "script_sha256": hashlib.sha256(Path(__file__).read_bytes()).hexdigest(),
        "requirements_sha256": hashlib.sha256((HERE/"frobenius-kerr-requirements.txt").read_bytes()).hexdigest(),
        "runtime": {"python": platform.python_version(), "machine": platform.machine(),
                    "packages": {p: importlib.metadata.version(p) for p in
                                 ("qnm", "numpy", "scipy", "numba", "llvmlite", "mpmath", "tqdm")}},
        "package_source_sha256": {x.__name__: hashlib.sha256(Path(x.__file__).read_bytes()).hexdigest()
                                  for x in modules},
        "scope": {"w": "GM*omega", "a_star": "a/(GM)", "time": "exp(-i*omega*t)",
                  "s": S, "ell": ELL, "m": M, "overtone": 0,
                  "continuation": "nine explicit coupled root solves, a_star=0..0.4",
                  "cached_or_interpolated_modes_used": False,
                  "package_import_note": "qnm initializes its bundled Schwarzschild dictionary and two warm-up solves; neither supplies this run's seeds/results",
                  "radial_independent_check": "Leaver1985 coefficients + 60dps backward zero-tail F0 at double precision w,A",
                  "angular_independent_check": "cos(theta) recurrence matrix, extended-space residual and ell_max re-solves",
                  "angular_series_check": "Leaver1985 Eqs18-21 minimal fraction at fixed double-precision w; not a new high-precision coupled root",
                  "input_contract": "no CLI arguments; only the declared one-mode finite benchmark is supported",
                  "claim_limit": "one low Kerr mode in double precision; diagnostics do not prove global mode completeness or certify 60-digit frequency accuracy"},
        "solve_controls": {"tol_xtol": 1e-11, "cf_tol": 1e-13,
                           "Nr_min": 200, "Nr_max": 8000, "ell_max": 12,
                           "initial_w_seed": [.37367, -.08896], "initial_A": [4, 0]},
        "continuation": path,
        "angular_resolution_coupled_resolves": resolution,
        "radial_independent_and_deeper_checks": radial,
        "angular_matrix_vs_series_checks": angular_series,
        "negative_controls": negative_controls,
        "independent_Leaver_vs_package_D_max_abs_difference": float(D_difference),
        "published_comparison": reference,
    }
    print(json.dumps(out, indent=2, sort_keys=True, allow_nan=False))


if __name__ == "__main__":
    main()