"""Checks on the exponential tail of numerics/tails.py requested in review.

Reports, for the 8,323 compounds within 50 meV/atom of the hull: the mean excess with
standard errors; an Anderson-Darling test of exponentiality above 10 K by parametric
bootstrap; the probability of the observed maximum under the fitted law; generalised
Pareto fits with the shape free and bootstrap intervals for the shape and return levels;
return levels for a constructed pool (p_u = 4.3%) and an unselected one (p_u = 0.27%);
and the dependence of the top of the sample on the position of the hull cut.
"""

import csv
import json
import math
from pathlib import Path

import numpy as np
from scipy.stats import genpareto

HERE = Path(__file__).parent
rows = [r for r in csv.reader(l for l in open(HERE.parent / "research/data/gao2025_fig4_tc_ehull.csv") if not l.startswith("#"))][1:]
tc = np.array([float(r[0]) for r in rows]); eh = np.array([float(r[1]) for r in rows])
rng = np.random.default_rng(3)
out = {}

near = tc[eh <= 50.0]
out["mean_excess"] = []
for u in (2, 5, 7.5, 10, 12.5, 15, 20, 25):
    e = near[near >= u] - u
    out["mean_excess"].append({"u": u, "n": int(e.size), "mean": float(e.mean()), "se": float(e.std(ddof=1) / math.sqrt(e.size))})

u = 10.0
exc = np.sort(near[near >= u] - u); n = exc.size; tau = exc.mean(); p_u = n / near.size


def ad_stat(x):
    x = np.sort(x); m = x.size; z = 1.0 - np.exp(-x / x.mean()); z = np.clip(z, 1e-12, 1 - 1e-12)
    i = np.arange(1, m + 1)
    return -m - np.sum((2 * i - 1) * (np.log(z) + np.log(1 - z[::-1]))) / m


a_obs = ad_stat(exc)
a_sim = np.array([ad_stat(rng.exponential(1.0, n)) for _ in range(4000)])
out["anderson_darling"] = {"statistic": float(a_obs), "p_value": float((a_sim >= a_obs).mean()), "n": int(n)}

m_obs = exc.max()
out["maximum"] = {"observed_K": float(m_obs + u), "expected_level_K": float(u + tau * math.log(n)),
                  "p_max_at_least_observed": float(1 - (1 - math.exp(-m_obs / tau)) ** n),
                  "second_highest_K": float(exc[-2] + u)}

near49 = tc[eh <= 49.0]; e49 = near49[near49 >= u] - u
out["cut_at_49"] = {"n": int(near49.size), "tau": float(e49.mean()), "max_K": float(near49.max()),
                    "p_max": float(1 - (1 - math.exp(-(near49.max() - u) / e49.mean())) ** e49.size)}


def gpd_fit(e):
    c, _, s = genpareto.fit(e, floc=0)
    return c, s


def gpd_level(c, s, npu, thr):
    return thr + (s / c) * (npu ** c - 1) if abs(c) > 1e-6 else thr + s * math.log(npu)


out["gpd"] = []
for thr in (10.0, 15.0):
    e = near[near >= thr] - thr; pu = e.size / near.size
    c, s = gpd_fit(e)
    boots = []
    for _ in range(1000):
        b = rng.choice(e, e.size)
        try:
            boots.append(gpd_fit(b))
        except Exception:
            pass
    cs = np.array([b[0] for b in boots])
    rec = {"threshold": thr, "n": int(e.size), "shape": float(c), "shape_lo": float(np.percentile(cs, 2.5)), "shape_hi": float(np.percentile(cs, 97.5)), "scale": float(s), "levels": {}}
    for k in (6, 10):
        lv = gpd_level(c, s, 10.0 ** k * pu, thr)
        bl = [gpd_level(bc, bs, 10.0 ** k * pu, thr) for bc, bs in boots]
        rec["levels"][f"1e{k}"] = {"level": float(lv), "lo": float(np.percentile(bl, 2.5)), "hi": float(np.percentile(bl, 97.5))}
    rec["p_above_300"] = float(pu * genpareto.sf(300 - thr, c, loc=0, scale=s))
    out["gpd"].append(rec)

out["exponential_levels"] = {}
for name, pu in (("constructed_sample_4.3pct", p_u), ("training_set_3.1pct", 0.031), ("unselected_pool_0.27pct", 0.0027)):
    out["exponential_levels"][name] = {f"1e{k}": float(u + tau * math.log(10.0 ** k * pu)) for k in (6, 8, 10, 12)}
out["sd_of_maximum_K"] = float(1.2825 * tau)
out["tau"] = float(tau); out["p_u"] = float(p_u)

(HERE / "tails_checks.json").write_text(json.dumps(out, indent=1))
for m in out["mean_excess"]:
    print(f"mean excess above {m['u']:>4} K: {m['mean']:.2f} +/- {m['se']:.2f} K (n={m['n']})")
print("Anderson-Darling:", out["anderson_darling"])
print("maximum:", out["maximum"])
print("cut at 49 meV:", out["cut_at_49"])
for g in out["gpd"]:
    print(f"GPD above {g['threshold']} K (n={g['n']}): shape {g['shape']:.2f} ({g['shape_lo']:.2f} to {g['shape_hi']:.2f}); 1e6 level {g['levels']['1e6']['level']:.0f} ({g['levels']['1e6']['lo']:.0f} to {g['levels']['1e6']['hi']:.0f}); 1e10 level {g['levels']['1e10']['level']:.0f} ({g['levels']['1e10']['lo']:.0f} to {g['levels']['1e10']['hi']:.0f}); P(>300) {g['p_above_300']:.1e}")
print("exponential levels:", {k: {a: round(b) for a, b in v.items()} for k, v in out["exponential_levels"].items()})
print("sd of maximum:", round(out["sd_of_maximum_K"], 1))