"""Multi-isotope survey: s-wave neutron resonance spacings vs Riemann zero spacings.
Run: python compare.py  (needs numpy, scipy; keep the adjacent data files).
Selection: exact duplicate energies merged. ENDF/B-VII.1, even-even targets (spin 0), first resolved region,
l=0, J=1/2, 0 < E <= upper bound of that region; isotopes with >= 100 levels.
KS thresholds and p-values are nominal (iid assumption); spacings are correlated,
unfolding is a fitted line, and evaluations miss weak levels. Diagnostic only.
"""
import csv, gzip, hashlib, json
from pathlib import Path
import numpy as np
from scipy.stats import ks_2samp
from scipy.special import erf

ROOT = Path(__file__).resolve().parent
ISOTOPES = ["n-092_U_238", "n-094_Pu_240", "n-090_Th_232", "n-064_Gd_154",
            "n-092_U_234", "n-068_Er_166", "n-068_Er_168"]


def endf_float(field):
    t = field.strip()
    i = max(t.rfind("+"), t.rfind("-"))
    return float(t[:i] + "e" + t[i:]) if i > 0 else float(t)


def fields(line):
    return [endf_float(line[i:i+11]) for i in range(0, 66, 11)]


def extract(raw):
    lines = [l for l in raw.splitlines() if l[70:72].strip() == "2" and l[72:75].strip() == "151"]
    region = fields(lines[2])
    el, eh, lru, lrf = region[0], region[1], int(region[2]), int(region[3])
    if lru != 1 or lrf not in (1, 2, 3) or int(region[4]) != 0:
        raise ValueError("unsupported resonance format")
    control = fields(lines[3])
    if control[0] != 0:
        raise ValueError("target spin must be 0")
    cursor, out = 4, []
    for _ in range(int(control[4])):
        h = fields(lines[cursor]); n = int(h[5])
        rows = [fields(l) for l in lines[cursor+1:cursor+1+n]]
        if len(rows) != n:
            raise ValueError("truncated block")
        if int(h[2]) == 0:
            out += [r[0] for r in rows if abs(r[1]) == 0.5 and 0 < r[0] <= eh and r[0] > el]
        cursor += 1 + n
    # Exact duplicate energies (one in Gd-154) are merged; count is recorded.
    return np.unique(out), el, eh, lrf, len(out) - len(np.unique(out))


def norm_spacing(x):
    d = np.diff(x)
    if np.any(d <= 0):
        raise ValueError("levels must be strictly increasing")
    return d / d.mean()


def ks(values, cdf):
    x = np.sort(values); n = len(x); up = np.arange(1, n+1) / n
    return float(np.max(np.maximum(abs(up-cdf(x)), abs(up-1/n-cdf(x)))))


GOE = lambda s: 1-np.exp(-np.pi*s*s/4)
GUE = lambda s: erf(2*s/np.sqrt(np.pi)) - 4*s/np.pi*np.exp(-4*s*s/np.pi)
POI = lambda s: 1-np.exp(-s)


def metrics(s):
    return {"spacings": len(s), "variance": float(s.var()),
            "small_gap_fraction": float(np.mean(s < 0.25)),
            "nominal_5pct_threshold": float(1.36/np.sqrt(len(s))),
            "ks_goe": ks(s, GOE), "ks_gue": ks(s, GUE), "ks_poisson": ks(s, POI)}


def verdict(m):
    t = m["nominal_5pct_threshold"]
    ok = [k for k in ("goe", "gue", "poisson") if m[f"ks_{k}"] < t]
    return "+".join(k.upper() for k in ok) if ok else "none"


def run(root=ROOT, write=True):
    zeros_all = np.loadtxt(root / "zeros-first2000.txt")
    results, rows = [], []
    for name in ISOTOPES:
        raw = gzip.decompress((root / f"{name}.mf2.endf.gz").read_bytes()).decode()
        e, el, eh, lrf, dups = extract(raw)
        fit = np.polyfit(e, np.arange(1, len(e)+1), 1)
        a = norm_spacing(np.polyval(fit, e))
        z = zeros_all[:len(e)]
        b = norm_spacing(z/(2*np.pi)*np.log(z/(2*np.pi*np.e)))
        c = ks_2samp(a, b, method="auto")
        ma, mb = metrics(a), metrics(b)
        results.append({
            "isotope": name.split("_", 1)[1].replace("_", "-"), "file": f"{name}.mf2.endf.gz",
            "sha256": hashlib.sha256(raw.encode()).hexdigest(), "lrf": lrf,
            "levels": len(e), "duplicates_merged": dups, "energy_max_eV": float(e[-1]), "region_upper_eV": eh,
            "mean_spacing_eV": float(1/fit[0]), "nucleus": ma, "zeros": mb,
            "nucleus_fits": verdict(ma), "zeros_fit": verdict(mb),
            "two_sample_ks": float(c.statistic), "nominal_two_sample_p": float(c.pvalue)})
        for i in range(len(a)):
            rows.append([results[-1]["isotope"], i+1, f"{e[i]:.6g}", f"{e[i+1]:.6g}",
                         f"{a[i]:.6f}", f"{z[i]:.9f}", f"{z[i+1]:.9f}", f"{b[i]:.6f}"])
    out = {"run_date": "2026-10-17", "library": "ENDF/B-VII.1 neutrons",
           "source_url": "https://www.nndc.bnl.gov/endf-b7.1/zips/ENDF-B-VII.1-neutrons.zip",
           "zero_source_url": "https://www-users.cse.umn.edu/~odlyzko/zeta_tables/zeros1",
           "caveat": "Nominal iid KS diagnostics; correlated spacings, fitted unfolding and missed weak levels are not calibrated. Rank pairing is a bookkeeping alignment, not a physical match.",
           "isotopes": results}
    if write:
        (root / "results.json").write_text(json.dumps(out, indent=2))
        with open(root / "rank-paired-spacings.csv", "w", newline="") as f:
            w = csv.writer(f)
            w.writerow(["isotope", "rank", "level_E_eV", "next_level_E_eV", "level_spacing_norm",
                        "zero_gamma", "next_zero_gamma", "zero_spacing_norm"])
            w.writerows(rows)
    return out


if __name__ == "__main__":
    for r in run()["isotopes"]:
        print(f'{r["isotope"]:7} n={r["levels"]:4} nucleus={r["nucleus_fits"]:8} zeros={r["zeros_fit"]:8} '
              f'KS2={r["two_sample_ks"]:.3f} p={r["nominal_two_sample_p"]:.2g} '
              f'goe={r["nucleus"]["ks_goe"]:.3f} gue={r["nucleus"]["ks_gue"]:.3f} thr={r["nucleus"]["nominal_5pct_threshold"]:.3f}')