"""Tail of calculated Tc in the ambient-pressure dataset of Gao et al. (2025).

Input: marker coordinates read from the vector figures in the arXiv source of that paper
(research/data/*.csv; Allen-Dynes Tc, harmonic, mu* = 0.1). For strata of energy above
the hull this fits the exponential law P(Tc >= T) = p_u exp(-(T - u) / tau) above a
threshold u, checks it with the mean excess at several thresholds, and reports the
return level u + tau ln(N p_u). Also reports the Pareto front of Tc against hull distance
and the largest lambda * omega_log^2 in the coupling-frequency data.
"""

import csv
import json
import math
from pathlib import Path

import numpy as np

HERE = Path(__file__).parent
DATA = HERE.parent / "research" / "data"
M_H_OMEGA2_1000K = 1.7904


def read(name):
    with open(DATA / name) as f:
        rows = [r for r in csv.reader(l for l in f if not l.startswith("#"))]
    head, body = rows[0], rows[1:]
    return {h: np.array([float(r[i]) for r in body]) for i, h in enumerate(head)}


def exp_fit(tc, u, rng, n_boot=2000):
    exc = tc[tc >= u] - u
    tau = exc.mean()
    boots = [rng.choice(exc, exc.size).mean() for _ in range(n_boot)]
    return {"u": u, "n": int(tc.size), "n_exceed": int(exc.size), "p_u": exc.size / tc.size,
            "tau": float(tau), "tau_lo": float(np.percentile(boots, 2.5)), "tau_hi": float(np.percentile(boots, 97.5)),
            "max": float(tc.max())}


def return_level(fit, n):
    return fit["u"] + fit["tau"] * math.log(n * fit["p_u"])


if __name__ == "__main__":
    rng = np.random.default_rng(7)
    d = read("gao2025_fig4_tc_ehull.csv")
    tc, eh = d["tc_allen_dynes_K"], d["e_hull_meV_per_atom"]
    out = {"n_points": int(tc.size)}

    near = tc[eh <= 50.0]
    out["near_hull_mean_excess"] = [
        {"u": u, "n_exceed": int((near >= u).sum()), "mean_excess": float((near[near >= u] - u).mean())}
        for u in (2, 5, 7.5, 10, 12.5, 15, 20)
    ]
    fit = exp_fit(near, 10.0, rng)
    out["near_hull_fit"] = fit
    out["near_hull_decade_K"] = fit["tau"] * math.log(10)
    out["near_hull_return_levels"] = {f"1e{k}": return_level(fit, 10.0**k) for k in (4, 6, 8, 10, 12)}
    out["near_hull_return_levels"]["sample"] = return_level(fit, near.size)
    out["near_hull_log10_N_for_300K"] = (300.0 - fit["u"]) / fit["tau"] / math.log(10) - math.log10(fit["p_u"])
    for lo, hi in ((fit["tau_lo"], "hi"), (fit["tau_hi"], "lo")):
        out[f"near_hull_log10_N_for_300K_{hi}"] = (300.0 - fit["u"]) / lo / math.log(10) - math.log10(fit["p_u"])
    out["near_hull_expected_above"] = {str(t): near.size * fit["p_u"] * math.exp(-(t - 10.0) / fit["tau"]) for t in (30, 39, 50)}
    out["near_hull_observed_above"] = {str(t): int((near >= t).sum()) for t in (30, 39, 50)}

    strata = [(0, 25), (25, 50), (50, 100), (100, 200), (200, 376)]
    out["strata"] = []
    for lo, hi in strata:
        sel = tc[(eh > lo if lo else eh >= 0) & (eh <= hi)]
        f = exp_fit(sel, 10.0, rng)
        f.update({"lo": lo, "hi": hi})
        out["strata"].append(f)

    # Kolmogorov-Smirnov distance of the exceedances from an exponential, near hull.
    exc = np.sort(near[near >= 10.0] - 10.0)
    cdf = 1.0 - np.exp(-exc / fit["tau"])
    emp = np.arange(1, exc.size + 1) / exc.size
    out["near_hull_ks"] = float(np.max(np.abs(cdf - emp)))

    # Pareto front: highest Tc at or below each hull distance.
    order = np.argsort(eh)
    best, front = -1.0, []
    for i in order:
        if tc[i] > best:
            best = tc[i]
            front.append({"e_hull": float(eh[i]), "tc": float(tc[i])})
    out["pareto_front"] = [p for p in front if p["tc"] >= 20.0]
    out["best_at_or_below"] = {str(t): float(tc[eh <= t].max()) for t in (0.5, 25, 50, 100, 200, 376)}
    out["count_at_or_below"] = {str(t): int((eh <= t).sum()) for t in (0.5, 25, 50, 100, 200, 376)}

    # survival curves for the figure
    grid = np.arange(0, 116, 1.0)
    out["survival"] = {"grid": grid.tolist()}
    for lo, hi in strata:
        sel = tc[(eh > lo if lo else eh >= 0) & (eh <= hi)]
        out["survival"][f"{lo}-{hi}"] = [float((sel >= g).mean()) for g in grid]

    g = read("gao2025_fig2_wlog_lambda_tc.csv")
    w, lam = g["omega_log_K"], g["lambda"]
    s_log = lam * (w / 1000.0) ** 2 * M_H_OMEGA2_1000K
    i = int(np.argmax(s_log))
    out["fig2"] = {
        "n": int(w.size), "max_omega_log": float(w.max()), "max_lambda": float(lam.max()),
        "max_S_lower_bound_eV_A2": float(s_log[i]), "at": {"omega_log": float(w[i]), "lambda": float(lam[i])},
        "n_S_above_1": int((s_log >= 1.0).sum()), "n_S_above_1p5": int((s_log >= 1.5).sum()),
        "n_lambda_ge2_and_wlog_ge500": int(((lam >= 2) & (w >= 500)).sum()),
        "n_wlog_ge_1000": int((w >= 1000).sum()),
        "asymptote_lower_bound_K": float(182.7 * math.sqrt(s_log[i] / M_H_OMEGA2_1000K)),
    }

    (HERE / "tails.json").write_text(json.dumps(out, indent=1))
    print("points", out["n_points"], "| within 50 meV/atom:", near.size)
    for m in out["near_hull_mean_excess"]:
        print(f"  mean excess above {m['u']:>4} K: {m['mean_excess']:.2f} K  ({m['n_exceed']} compounds)")
    print(f"fit u=10 K: p_u={fit['p_u']:.4f}, tau={fit['tau']:.2f} K ({fit['tau_lo']:.2f} to {fit['tau_hi']:.2f}); one decade = {out['near_hull_decade_K']:.1f} K; KS distance {out['near_hull_ks']:.3f} with n={fit['n_exceed']}")
    print("return levels:", {k: round(v, 1) for k, v in out["near_hull_return_levels"].items()})
    print(f"log10 N for 300 K: {out['near_hull_log10_N_for_300K']:.1f} ({out['near_hull_log10_N_for_300K_lo']:.1f} to {out['near_hull_log10_N_for_300K_hi']:.1f})")
    print("expected above 30/39/50:", {k: round(v, 3) for k, v in out["near_hull_expected_above"].items()}, "observed:", out["near_hull_observed_above"])
    for s in out["strata"]:
        print(f"  {s['lo']:>3}-{s['hi']:<3} meV/atom: n={s['n']:5d}, above 10 K {s['n_exceed']:4d}, tau={s['tau']:5.2f} ({s['tau_lo']:.2f} to {s['tau_hi']:.2f}), max {s['max']:.1f} K")
    print("best Tc at or below hull distance:", out["best_at_or_below"])
    print("front:", [(round(p['e_hull'], 1), round(p['tc'], 1)) for p in out["pareto_front"]])
    print("fig2:", out["fig2"])