"""Census of the scattering strength per proton, h = eta_H / rho_H, across hydrides.

Reads data/h_census.json. Each row carries published inputs (lambda, omega_2, the number of
H atoms and the volume of the primitive cell, and eta_H where a paper prints it). This
script recomputes

    S     = M_H omega^2(1000 K) * lambda * (omega_2 / 1000 K)^2      eV/A^2
    rho_H = n_H / V                                                  H per A^3
    h     = eta_H / rho_H  if eta_H is published, else  S / rho_H    eV A

checks them against the values stored in the file, recomputes the Allen-Dynes Tc of every
row as a guard against transcription slips, and prints the statistics that the paper quotes.
The statistics are also written to numerics/h_census_stats.json.

S is the Hopfield sum in hydrogen units. It equals eta_H plus eta_X M_H / M_X for every
other atom, so it overstates eta_H where a light atom couples or where hydrogen is a small
part of the cell. The file flags both cases.

Usage:  .venv/bin/python numerics/h_census.py
"""

import json
import math
from pathlib import Path

import numpy as np
from scipy import stats

HERE = Path(__file__).parent
DATA = HERE.parent / "data" / "h_census.json"
OUT = HERE / "h_census_stats.json"

AMU = 1.66053906660e-27
KB = 1.380649e-23
HBAR = 1.054571817e-34
EV = 1.602176634e-19
TOL = 1.5e-3  # stored values are rounded to 4 or 5 significant figures


def m_h_omega2_1000k(m_h_u):
    omega = KB * 1000.0 / HBAR
    return m_h_u * AMU * omega**2 / EV * 1e-20


def recompute(row, c):
    lam, w2 = row.get("lambda"), row.get("omega_2_K")
    s = c["M_H_omega2_at_1000K_eV_A2"] * lam * (w2 / 1000.0) ** 2 if lam is not None and w2 is not None else None
    rho = row["n_H"] / row["volume_A3"] if row.get("volume_A3") else None
    eta = row.get("eta_H_eV_A2")
    num = eta if eta is not None else s
    h = num / rho if num is not None and rho else None
    h_s = s / rho if s is not None and rho else None
    return s, rho, h, h_s


def allen_dynes(lam, wlog, w2, mu):
    """Allen-Dynes Tc with the strong-coupling and shape factors f1, f2."""
    denom = lam - mu * (1.0 + 0.62 * lam)
    if denom <= 0:
        return 0.0
    l1 = 2.46 * (1.0 + 3.8 * mu)
    l2 = 1.82 * (1.0 + 6.3 * mu) * (w2 / wlog)
    f1 = (1.0 + (lam / l1) ** 1.5) ** (1.0 / 3.0)
    f2 = 1.0 + (w2 / wlog - 1.0) * lam**2 / (lam**2 + l2**2)
    return f1 * f2 * wlog / 1.2 * math.exp(-1.04 * (1.0 + lam) / denom)


def summary(values):
    x = np.asarray([v for v in values if v is not None and v > 0], dtype=float)
    if len(x) == 0:
        return {"n": 0}
    ln = np.log(x)
    q25, q50, q75 = np.percentile(x, [25, 50, 75])
    sd = float(ln.std(ddof=1)) if len(x) > 1 else 0.0
    return {
        "n": int(len(x)),
        "geometric_mean": float(np.exp(ln.mean())),
        "sd_ln": sd,
        "geometric_sd_factor": float(math.exp(sd)),
        "median": float(q50),
        "min": float(x.min()),
        "max": float(x.max()),
        "q25": float(q25),
        "q75": float(q75),
        "iqr": float(q75 - q25),
        "max_over_min": float(x.max() / x.min()),
    }


def corr(x, y):
    pairs = [(a, b) for a, b in zip(x, y) if a is not None and b is not None]
    if len(pairs) < 4:
        return {"n": len(pairs)}
    a, b = np.array(pairs, dtype=float).T
    r, p = stats.pearsonr(a, b)
    rs, ps = stats.spearmanr(a, b)
    return {"n": len(pairs), "pearson_r": float(r), "pearson_p": float(p), "spearman_rho": float(rs), "spearman_p": float(ps)}


def slope(x, y):
    fit = stats.linregress(x, y)
    return {"n": len(x), "slope": float(fit.slope), "slope_se": float(fit.stderr), "intercept": float(fit.intercept), "r": float(fit.rvalue)}


def fmt(s, unit=""):
    if s["n"] == 0:
        return "n = 0"
    return (
        f"n = {s['n']:2d}  gmean {s['geometric_mean']:7.2f}{unit}  sd(ln) {s['sd_ln']:.3f} (x{s['geometric_sd_factor']:.2f})  "
        f"median {s['median']:7.2f}  range {s['min']:.2f} to {s['max']:.2f}  IQR {s['q25']:.2f} to {s['q75']:.2f} (width {s['iqr']:.2f})"
    )


def main():
    data = json.loads(DATA.read_text())
    c = data["constants"]
    rows = data["rows"]
    two_e2 = 2.0 * c["e2_eV_A"]

    exact = m_h_omega2_1000k(c["M_H_u"])
    assert abs(exact - c["M_H_omega2_at_1000K_eV_A2"]) < 1e-4, exact
    print(f"M_H omega^2 at 1000 K = {exact:.5f} eV/A^2 for M_H = {c['M_H_u']} u (file uses {c['M_H_omega2_at_1000K_eV_A2']});  2e^2 = {two_e2:.4f} eV A")

    # 1. derived fields
    worst = 0.0
    for r in rows:
        s, rho, h, h_s = recompute(r, c)
        for key, val in (("S_eV_A2", s), ("rho_H", rho), ("h_eV_A", h), ("h_from_S_eV_A", h_s)):
            stored = r.get(key)
            if (val is None) != (stored is None):
                raise SystemExit(f"{r['id']}: {key} stored {stored}, recomputed {val}")
            if val is not None:
                dev = abs(val - stored) / abs(val)
                worst = max(worst, dev)
                if dev > TOL:
                    raise SystemExit(f"{r['id']}: {key} stored {stored}, recomputed {val:.5f}")
        if h is not None and abs(h / two_e2 - r["h_over_2e2"]) > 2e-3:
            raise SystemExit(f"{r['id']}: h_over_2e2 stored {r['h_over_2e2']}, recomputed {h / two_e2:.4f}")
        r["_S"], r["_rho"], r["_h"], r["_hS"] = s, rho, h, h_s
    print(f"derived fields recomputed for {len(rows)} rows; largest relative difference from the stored value {worst:.1e}")

    # 2. Allen-Dynes check of the transcribed lambda, omega_log, omega_2 against the printed Tc
    ad = []
    for r in rows:
        if r["pressure_GPa"] != 0 or r.get("omega_2_K") is None or r.get("tc_K") is None or "Allen-Dynes" not in (r.get("tc_method") or ""):
            continue
        t = allen_dynes(r["lambda"], r["omega_log_K"], r["omega_2_K"], 0.10)
        ad.append((r["id"], r["compound"], r["tc_K"], t))
    dev = [abs(t - p) for _, _, p, t in ad]
    print(f"Allen-Dynes (mu* = 0.10) recomputed for {len(ad)} rows at 1 atm: largest |difference| from the printed Tc {max(dev):.2f} K, mean {np.mean(dev):.2f} K")
    for i, comp, p, t in ad:
        if abs(t - p) > max(1.0, 0.03 * p):
            print(f"    check {i} {comp}: printed {p} K, recomputed {t:.1f} K")
    ad_mega = []
    for r in rows:
        if r["group"] == "megabar" and r.get("omega_2_K") and r.get("tc_K"):
            ad_mega.append((r["id"], r["tc_K"], allen_dynes(r["lambda"], r["omega_log_K"], r["omega_2_K"], 0.13)))
    dm = [abs(t - p) / p for _, p, t in ad_mega]
    print(f"Allen-Dynes (mu* = 0.13) for {len(ad_mega)} megabar total rows: largest relative difference from the printed Tc {max(dm):.3f}")

    # 3. groups
    counted = [r for r in rows if r["include_in_stats"]]
    amb = [r for r in counted if r["group"] == "ambient"]
    exi = [r for r in counted if r["group"] == "ambient-existing"]
    meg = [r for r in counted if r["group"] == "megabar"]
    reviewer = {x["compound"]: x for x in data["reviewer_table"]["rows"]}
    rev12 = [r for r in amb if r["compound"] in reviewer and r["id"].startswith(("cerqueira", "gao-Li2AgH6"))]
    groups = {
        "a": ("(a) ambient, all counted rows", amb),
        "a_cerqueira": ("    of which Cerqueira Tables I-II (Tc above 20 K)", [r for r in amb if r["id"].startswith("cerqueira")]),
        "a_firm": ("    of which structure match not 'formula_only'", [r for r in amb if r["structure_match"] != "formula_only"]),
        "b": ("(b) ambient, no light partner", [r for r in amb if not r["light_partner"]]),
        "b_hrich": ("    of which H is at least half the atoms", [r for r in amb if not r["light_partner"] and r["h_atom_fraction"] >= 0.5]),
        "b_light": ("    ambient rows with a light partner", [r for r in amb if r["light_partner"]]),
        "reviewer12": ("    the twelve compounds of the reviewer's table", rev12),
        "c": ("(c) megabar, all pressures", meg),
        "c_one": ("    one pressure per compound (lowest)", [r for r in meg if r.get("first_pressure_of_compound")]),
        "d": ("(d) everything counted", amb + exi + meg),
        "d_one": ("    megabar at one pressure per compound", amb + exi + [r for r in meg if r.get("first_pressure_of_compound")]),
    }
    out = {
        "constants": {"two_e2_eV_A": two_e2, "M_H_omega2_at_1000K_eV_A2": c["M_H_omega2_at_1000K_eV_A2"]},
        "row_counts": {},
        "groups": {},
    }
    for g in ("ambient", "ambient-existing", "megabar"):
        out["row_counts"][g] = {
            "rows": sum(r["group"] == g for r in rows),
            "counted": sum(r["group"] == g and r["include_in_stats"] for r in rows),
            "with_h": sum(r["group"] == g and r["_h"] is not None for r in rows),
        }
    print("\nrows per group (total, counted in statistics):", {g: (v["rows"], v["counted"]) for g, v in out["row_counts"].items()})
    print("megabar h uses the published eta_H; ambient h uses S. Megabar rows at several pressures of one compound are not independent.")

    for quantity, key, unit in (("h (eV A)", "_h", ""), ("S (eV/A^2)", "_S", ""), ("Tc (K)", "tc_K", "")):
        print(f"\n{quantity}")
        for gk, (label, rs) in groups.items():
            s = summary([r.get(key) for r in rs])
            out["groups"].setdefault(gk, {"label": label.strip(), "n_rows": len(rs)})[key.strip("_")] = s
            print(f"  {label:52s} {fmt(s, unit)}")
    s_megabar_from_S = summary([r["_hS"] for r in meg])
    out["groups"]["c"]["h_from_S"] = s_megabar_from_S
    print(f"\n  megabar h recomputed from S instead of eta_H         {fmt(s_megabar_from_S)}")
    zero_tc = [r["compound"] for r in amb if r.get("tc_K") == 0]
    print(f"  rows with Tc = 0 left out of the Tc statistics: {zero_tc}")

    # 4. width of h against S and Tc
    print("\nsd of the logarithm: h against S and Tc")
    out["width_ratios"] = {}
    for gk in ("a", "a_cerqueira", "b", "b_hrich", "c", "d"):
        g = out["groups"][gk]
        rs_ = g["S"]["sd_ln"] / g["h"]["sd_ln"]
        rt = g["tc_K"]["sd_ln"] / g["h"]["sd_ln"]
        out["width_ratios"][gk] = {"sd_ln_h": g["h"]["sd_ln"], "sd_ln_S": g["S"]["sd_ln"], "sd_ln_Tc": g["tc_K"]["sd_ln"], "S_over_h": rs_, "Tc_over_h": rt}
        print(f"  {gk:12s} sd ln h {g['h']['sd_ln']:.3f}   sd ln S {g['S']['sd_ln']:.3f} ({rs_:.2f} x)   sd ln Tc {g['tc_K']['sd_ln']:.3f} ({rt:.2f} x)")

    # 5. how many ambient rows a constant reproduces
    print("\nambient rows within 30% of a constant (|h / constant - 1| <= 0.30)")
    out["within_30pct"] = {}
    for gk in ("a", "a_cerqueira", "b", "b_hrich", "reviewer12"):
        hs = np.array([r["_h"] for r in groups[gk][1]])
        gm = out["groups"][gk]["h"]["geometric_mean"]
        n2 = int(np.sum(np.abs(hs / two_e2 - 1.0) <= 0.30))
        ng = int(np.sum(np.abs(hs / gm - 1.0) <= 0.30))
        npred = int(np.sum(np.abs(two_e2 / hs - 1.0) <= 0.30))
        out["within_30pct"][gk] = {
            "n": int(len(hs)),
            "of_2e2": n2,
            "of_group_geometric_mean": ng,
            "group_geometric_mean": gm,
            "2e2_as_prediction_error_within_30pct": npred,
            "below_2e2": int(np.sum(hs < two_e2)),
        }
        print(f"  {gk:12s} n = {len(hs):2d}: {n2} within 30% of 2e^2, {ng} within 30% of the group geometric mean {gm:.1f}; "
              f"{npred} if 2e^2 is scored as a prediction of h (reviewer's convention); {int(np.sum(hs < two_e2))} lie below 2e^2")

    # 6. correlations
    print("\ncorrelations (Pearson r, Spearman rho)")
    out["correlations"] = {}
    for gk in ("a", "a_cerqueira", "b", "b_hrich", "d"):
        rs = groups[gk][1]
        lnh = [math.log(r["_h"]) for r in rs]
        lnrho = [math.log(r["_rho"]) for r in rs]
        lnS = [math.log(r["_S"]) if r["_S"] else None for r in rs]
        res = {
            "ln_h_vs_ln_rho_H": corr(lnh, lnrho),
            "ln_h_vs_lambda": corr(lnh, [r.get("lambda") for r in rs]),
            "ln_h_vs_e_hull": corr(lnh, [r.get("e_hull_meV") for r in rs]),
            "ln_h_vs_ln_tc": corr(lnh, [math.log(r["tc_K"]) if r.get("tc_K") else None for r in rs]),
            "ln_S_vs_ln_rho_H": corr(lnS, lnrho),
        }
        ok = [(a, b) for a, b in zip(lnrho, lnS) if b is not None]
        res["slope_ln_S_on_ln_rho_H"] = slope([a for a, _ in ok], [b for _, b in ok])
        v_h, v_rho, v_s = np.var(lnh, ddof=1), np.var(lnrho, ddof=1), np.var([b for _, b in ok], ddof=1)
        res["variances"] = {"ln_h": float(v_h), "ln_rho_H": float(v_rho), "ln_S": float(v_s)}
        out["correlations"][gk] = res
        print(f"  {groups[gk][0].strip()}")
        for name in ("ln_h_vs_ln_rho_H", "ln_h_vs_lambda", "ln_h_vs_e_hull", "ln_h_vs_ln_tc", "ln_S_vs_ln_rho_H"):
            q = res[name]
            if "pearson_r" in q:
                print(f"      {name:20s} n = {q['n']:2d}  r = {q['pearson_r']:+.2f} (p = {q['pearson_p']:.3f})  rho = {q['spearman_rho']:+.2f} (p = {q['spearman_p']:.3f})")
        sl = res["slope_ln_S_on_ln_rho_H"]
        print(f"      slope of ln S on ln rho_H: {sl['slope']:.2f} +/- {sl['slope_se']:.2f} (1 if h were the same in every compound, 0 if S ignored density)")
        print(f"      var ln S {v_s:.3f}, var ln rho_H {v_rho:.3f}, var ln h {v_h:.3f}")

    # 6b. h against the quantity the sample was selected on
    print("\nambient h by printed Tc (the tables were cut at Tc = 20 K)")
    out["h_by_tc_bin"] = []
    for lo, hi in ((0.0, 30.0), (30.0, 50.0), (50.0, 1e9)):
        sel = [r["_h"] for r in amb if r.get("tc_K") is not None and lo <= r["tc_K"] < hi]
        q = summary(sel)
        out["h_by_tc_bin"].append({"tc_from_K": lo, "tc_below_K": hi if hi < 1e9 else None, "h": q})
        top = f"{hi:.0f}" if hi < 1e9 else "up"
        print(f"  Tc {lo:3.0f} to {top:>3s} K   {fmt(q)}")

    # 7. extremes
    ranked = sorted(amb + exi + meg, key=lambda r: r["_h"])
    out["smallest_h"] = [{"id": r["id"], "compound": r["compound"], "group": r["group"], "pressure_GPa": r["pressure_GPa"], "h_eV_A": r["_h"], "light_partner": r["light_partner"], "h_atom_fraction": r["h_atom_fraction"]} for r in ranked[:5]]
    out["largest_h"] = [{"id": r["id"], "compound": r["compound"], "group": r["group"], "pressure_GPa": r["pressure_GPa"], "h_eV_A": r["_h"], "light_partner": r["light_partner"], "h_atom_fraction": r["h_atom_fraction"]} for r in ranked[::-1][:5]]
    print("\nfive smallest h among counted rows:")
    for r in ranked[:5]:
        print(f"    {r['compound']:10s} {r['_h']:6.2f}  ({r['group']}, {r['pressure_GPa']} GPa)")
    print("five largest h among counted rows:")
    for r in ranked[::-1][:5]:
        tag = "light partner" if r["light_partner"] else ""
        print(f"    {r['compound']:10s} {r['_h']:6.2f}  ({r['group']}, {r['pressure_GPa']} GPa) {tag}")

    # 8. comparison with the reviewer's table
    print("\ncomparison with reviews/ml-researcher.md (h recomputed here, reviewer's h, difference)")
    out["reviewer_comparison"] = []
    for r in rows:
        if r["compound"] in reviewer and r["id"].startswith(("cerqueira", "gao-Li2AgH6")):
            hr = reviewer[r["compound"]]["h"]
            d = r["_h"] / hr - 1.0
            out["reviewer_comparison"].append({"compound": r["compound"], "h_here": r["_h"], "h_reviewer": hr, "relative_difference": d})
            flag = "  <-- differs" if abs(d) > 0.01 else ""
            print(f"    {r['compound']:10s} {r['_h']:6.2f}  {hr:6.1f}  {100 * d:+.1f}%{flag}")
    g = out["groups"]["reviewer12"]["h"]
    w = out["within_30pct"]["reviewer12"]
    print(f"    twelve without PdH: geometric mean {g['geometric_mean']:.1f} (reviewer 26.4), sd ln h {g['sd_ln']:.3f} (0.228), "
          f"range {g['min']:.1f} to {g['max']:.1f} (17.4 to 37.3), {w['2e2_as_prediction_error_within_30pct']} of 12 within 30% (10)")
    for m in data["reviewer_table"]["megabar"]:
        n_h = {"LaH10": 10, "CaH6": 6}[m["compound"]]
        h_rev = m["eta_H"] / (n_h / m["volume_A3"])
        here = next(r for r in meg if r["compound"] == m["compound"] and r["pressure_GPa"] == m["pressure_GPa"])
        out["reviewer_comparison"].append({"compound": m["compound"], "pressure_GPa": m["pressure_GPa"], "h_here": here["_h"], "h_reviewer": h_rev,
                                           "relative_difference": here["_h"] / h_rev - 1.0, "volume_here_A3": here["volume_A3"], "volume_reviewer_A3": m["volume_A3"]})
        print(f"    {m['compound']} at {m['pressure_GPa']} GPa: {here['_h']:.1f} here (V = {here['volume_A3']} A^3), {h_rev:.1f} from the reviewer's volume ({m['volume_A3']} A^3), {100 * (here['_h'] / h_rev - 1):+.0f}%  <-- differs")

    # 9. functional of the cell
    chk = data["pbesol_pbe_volume_check"]
    ratios = np.array([x["ratio"] for x in chk["rows"]])
    core = ratios[(ratios > 0.92) & (ratios < 1.0)]
    out["pbesol_over_pbe_volume"] = {"n": int(len(ratios)), "median": float(np.median(ratios)), "n_between_0.92_and_1": int(len(core)),
                                     "min_core": float(core.min()), "max_core": float(core.max()), "median_core": float(np.median(core))}
    print(f"\ncell functional: PBEsol / PBE volume for {len(ratios)} hydride records, median {np.median(ratios):.3f}; "
          f"{len(core)} lie between {core.min():.3f} and {core.max():.3f} (median {np.median(core):.3f}).")
    shift = float(np.median(core))
    print(f"    ambient cells are PBE, the coupling is PBEsol: with PBEsol cells every ambient h would be about {100 * (1 - shift):.0f}% lower "
          f"(geometric mean of (a) {out['groups']['a']['h']['geometric_mean'] * shift:.1f} instead of {out['groups']['a']['h']['geometric_mean']:.1f}); sd(ln h) changes only through the scatter of that ratio.")
    out["groups"]["a"]["h_geometric_mean_with_pbesol_cells_estimate"] = out["groups"]["a"]["h"]["geometric_mean"] * shift

    # 10. rows kept out of the statistics
    print("\nrows listed but not counted:")
    out["not_counted"] = []
    for r in rows:
        if not r["include_in_stats"]:
            hval = f"{r['_h']:.1f}" if r["_h"] is not None else "none"
            out["not_counted"].append({"id": r["id"], "compound": r["compound"], "variant": r.get("variant"), "h_eV_A": r["_h"], "reason": r["excluded_because"]})
            print(f"    {r['id']:30s} {r['compound']:8s} h = {hval:>5s}  {r['excluded_because']}")

    out["selection"] = (
        "Ambient rows come from tables of compounds whose harmonic Allen-Dynes Tc exceeded 20 K after a machine-learning "
        "pre-screen (Cerqueira et al.), plus compounds singled out in Gao et al. for a high Tc or a large lambda. A cut on Tc "
        "removes weak coupling, so low values of h are missing. Two counted rows were not selected on Tc: NaNiH3 (lambda 0.16, "
        "h = 7.4) and Tl2AgH2 (Tc 12.7 K, h = 20.6). PdH passes the cut on a harmonic Tc that anharmonicity removes and has h = 7.2."
    )
    print("\nselection: " + out["selection"])
    OUT.write_text(json.dumps(out, indent=1) + "\n")
    print(f"\nwrote {OUT.relative_to(HERE.parent)}")


if __name__ == "__main__":
    main()