"""Census of the scattering strength per proton over the Alexandria electron-phonon release.

Input:  model/data/labels.csv and model/data/summary.json (written by model/dataset.py from
        the release of 11 August 2025), and data/h_census.json (the census of published
        tables used in draft 4 of the paper).
Output: numerics/h_alexandria.json

A hydride enters the statistics of h if every one of its 3 n_H highest phonon branches has
at least 90% of its eigenvector weight on hydrogen at every stored wavevector. For those,
eta_H is the part of the Hopfield sum in these branches and h = eta_H / rho_H. Labels are
at a smearing of 0.030 Ry unless a width is named.

Run:  .venv/bin/python numerics/h_alexandria.py
"""
import csv
import json
import math
import pathlib

import numpy as np
from pymatgen.core import Composition

ROOT = pathlib.Path(__file__).resolve().parents[1]
GAS = 36.6            # eV A, one proton in an electron gas at r_s = 2.5
K = 136.54            # K per sqrt(eV/A^2): asymptote of Tc from the Hopfield sum
FLOOR = 4.83          # eV/A^2: eta_H below which 300 K is impossible at any efficiency


def num(x):
    return float(x) if x not in ('', None) else float('nan')


def geo(v):
    lv = np.log(v)
    return dict(n=int(len(v)), geometric_mean=float(np.exp(lv.mean())), spread_factor=float(np.exp(lv.std())),
                quantiles={str(q): float(np.percentile(v, q)) for q in (5, 10, 25, 50, 75, 90, 95, 99)},
                min=float(v.min()), max=float(v.max()))


def line(x, y):
    a = np.vstack([x, np.ones_like(x)]).T
    coef = np.linalg.lstsq(a, y, rcond=None)[0]
    res = y - a @ coef
    se = math.sqrt(res.var(ddof=2) / ((x - x.mean()) ** 2).sum())
    return dict(slope=float(coef[0]), se=float(se), intercept=float(coef[1]), r2=float(1 - res.var() / y.var()), n=int(len(x)))


def main():
    rows = list(csv.DictReader(open(ROOT / 'model' / 'data' / 'labels.csv')))
    summary = json.load(open(ROOT / 'model' / 'data' / 'summary.json'))
    hyd = [r for r in rows if int(r['n_H']) > 0]
    lab = [r for r in hyd if r['branch_ok'] == 'True' and num(r['h']) > 0]
    h = np.array([num(r['h']) for r in lab])
    rho = np.array([num(r['rho_H']) for r in lab])
    eta = h * rho
    tc = np.array([num(r['tc_ad']) for r in lab])
    tc = np.where(np.isfinite(tc), tc, 0.0)
    out = dict(source='Alexandria phonon and electron-phonon release of 2025-08-11 (PBEsol, 3D)', counts=summary,
               labelled_hydrides=len(lab), smearing_Ry=0.030)

    out['h'] = geo(h)
    out['h']['fraction_above_electron_gas_value'] = float((h > GAS).mean())
    out['h']['count_above_electron_gas_value'] = int((h > GAS).sum())
    out['h']['fraction_above_22'] = float((h > 22).mean())
    out['h_by_Tc'] = {f'{lo} to {hi} K': (geo(h[(tc >= lo) & (tc < hi)]) if ((tc >= lo) & (tc < hi)).sum() > 2 else None)
                      for lo, hi in ((0, 1), (1, 5), (5, 10), (10, 20), (20, 1000))}
    out['h_by_density'] = {f'{lo} to {hi}': (geo(h[(rho >= lo) & (rho < hi)]) if ((rho >= lo) & (rho < hi)).sum() > 2 else None)
                           for lo, hi in ((0, 0.02), (0.02, 0.04), (0.04, 0.06), (0.06, 0.08), (0.08, 0.10), (0.10, 1.0))}
    out['eta_on_density'] = line(np.log(rho), np.log(eta))
    out['correlation_ln_h_ln_density'] = float(np.corrcoef(np.log(h), np.log(rho))[0, 1])

    out['eta_H'] = geo(eta)
    out['eta_H']['count_above'] = {str(t): int((eta > t).sum()) for t in (1, 2, 3, FLOOR, 8.7)}
    out['rho_H'] = geo(rho)
    out['rho_H']['count_above'] = {str(t): int((rho > t).sum()) for t in (0.08, 0.09, 0.10, 0.12)}
    share = np.array([num(r['S_top']) / num(r['S']) for r in lab])
    out['hydrogen_share_of_S'] = {str(q): float(np.percentile(share, q)) for q in (10, 25, 50, 75, 90)}

    def table(order, n=15):
        return [dict(formula=lab[i]['formula'], id=lab[i]['id'], spacegroup=int(lab[i]['spacegroup']),
                     rho_H=round(float(rho[i]), 4), eta_H=round(float(eta[i]), 3), h=round(float(h[i]), 1),
                     lam=round(num(lab[i]['lam']), 2), tc_allen_dynes=round(float(tc[i]), 1)) for i in order[:n]]
    out['largest_eta_H'] = table(np.argsort(-eta))
    out['largest_h'] = table(np.argsort(-h))
    out['largest_rho_H'] = table(np.argsort(-rho))

    # every hydride with a spectrum, whether or not the hydrogen branches separate: the whole sum
    s_all = np.array([num(r['S']) for r in hyd])
    rho_all = np.array([num(r['rho_H']) for r in hyd])
    out['all_hydrides_with_spectrum'] = dict(n=len(hyd), S=geo(s_all), S_over_rho_H=geo(s_all / rho_all),
                                             count_S_above={str(t): int((s_all > t).sum()) for t in (2, 3, FLOOR, 8.7)},
                                             largest_S=[dict(formula=hyd[i]['formula'], S=round(float(s_all[i]), 3), rho_H=round(float(rho_all[i]), 4),
                                                             branch_ok=hyd[i]['branch_ok'] == 'True') for i in np.argsort(-s_all)[:10]])

    # smearing
    widths = {}
    for tag in ('005', '050'):
        v = np.array([num(r['S_top_' + tag]) for r in lab]) / rho
        ok = np.isfinite(v) & (v > 0)
        widths['0.' + tag] = geo(v[ok])
    widths['0.030'] = dict(n=len(h), geometric_mean=out['h']['geometric_mean'], spread_factor=out['h']['spread_factor'])
    ls = np.array([math.log(num(r['S_005']) / num(r['S'])) for r in rows if num(r['S_005']) > 0 and num(r['S']) > 0])
    ll = np.array([math.log(num(r['lam_005']) / num(r['lam'])) for r in rows if num(r['lam_005']) > 0 and num(r['lam']) > 0])
    out['smearing'] = dict(h_by_width=widths,
                           ln_ratio_0005_over_0030=dict(S=dict(mean=float(ls.mean()), sd=float(ls.std()), n=int(len(ls))),
                                                        lam=dict(mean=float(ll.mean()), sd=float(ll.std()), n=int(len(ll)))))

    # the Tc-selected census of draft 4, matched by reduced formula and space group where given
    census = json.load(open(ROOT / 'data' / 'h_census.json'))['rows']
    by_formula = {}
    for r in hyd:
        by_formula.setdefault(Composition(r['formula']).reduced_formula, []).append(r)
    matched = []
    for c in census:
        if c.get('group') not in ('ambient', 'ambient-existing') or not c.get('S_eV_A2'):
            continue
        try:
            key = Composition(c['compound']).reduced_formula
        except Exception:
            continue
        cand = by_formula.get(key, [])
        if cand:
            best = min(cand, key=lambda r: abs(math.log(num(r['S']) / c['S_eV_A2'])))
            matched.append(dict(compound=c['compound'], S_census=c['S_eV_A2'], S_release=round(num(best['S']), 3),
                                ln_ratio=round(math.log(num(best['S']) / c['S_eV_A2']), 3), candidates=len(cand)))
    ratios = np.array([m['ln_ratio'] for m in matched])
    out['census_of_draft_4'] = dict(matched=len(matched), median_abs_ln_ratio=float(np.median(np.abs(ratios))) if len(ratios) else None,
                                    within_10pct=int((np.abs(ratios) < math.log(1.1)).sum()) if len(ratios) else 0, rows=matched)

    # what these strengths give at a hydrogen density of 0.10 per A^3
    def at(hv, phi):
        return K * math.sqrt(0.10 * hv) * phi
    out['at_density_0.10'] = {name: {'h': float(v), 'asymptote_K': at(v, 1.0), 'Tc_at_efficiency_0.35': at(v, 0.35), 'Tc_at_efficiency_0.58': at(v, 0.58)}
                              for name, v in (('median', np.percentile(h, 50)), ('geometric_mean', out['h']['geometric_mean']),
                                              ('90th_percentile', np.percentile(h, 90)), ('99th_percentile', np.percentile(h, 99)),
                                              ('largest', h.max()), ('electron_gas', GAS))}
    out['needed_for_300K'] = dict(eta_H_single_mode=[8.7, 12.0], eta_H_floor=FLOOR,
                                  ratio_of_single_mode_requirement_to_largest_eta_H=[8.7 / eta.max(), 12.0 / eta.max()],
                                  ratio_of_floor_to_largest_eta_H=FLOOR / eta.max())

    # upper tail of eta_H: mean excess, and an exponential fit to the top tenth
    u = float(np.percentile(eta, 90))
    exc = eta[eta > u] - u
    tau = float(exc.mean())
    boot = [np.random.default_rng(i).choice(exc, len(exc)).mean() for i in range(2000)]
    out['eta_H_tail'] = dict(threshold=u, n_above=int(len(exc)), mean_excess=tau, mean_excess_interval=[float(np.percentile(boot, 2.5)), float(np.percentile(boot, 97.5))],
                             mean_excess_by_threshold={str(q): (float((eta[eta > np.percentile(eta, q)] - np.percentile(eta, q)).mean()), int((eta > np.percentile(eta, q)).sum())) for q in (50, 75, 90, 95)},
                             exponential_fraction_above={str(t): float(0.1 * math.exp(-(t - u) / tau)) for t in (FLOOR, 8.7)})

    # coverage: how many hydride records carry a label at all
    out['coverage'] = dict(hydride_records=summary['hydride_records'], hydrides_with_spectrum=len(hyd), labelled=len(lab),
                           share_of_hydride_records_with_spectrum=len(hyd) / summary['hydride_records'],
                           share_of_all_records_with_imaginary_modes=summary['records_with_imaginary_modes'] / summary['records'])

    # upper edge of h against hydrogen density
    edges = [0, 0.02, 0.04, 0.06, 0.08, 0.10, 0.12, 1.0]
    out['h_upper_edge_by_density'] = []
    for lo, hi in zip(edges[:-1], edges[1:]):
        m = (rho >= lo) & (rho < hi)
        if m.sum() > 2:
            out['h_upper_edge_by_density'].append(dict(rho_H=[lo, hi], n=int(m.sum()), median=float(np.percentile(h[m], 50)),
                                                       p90=float(np.percentile(h[m], 90)), p99=float(np.percentile(h[m], 99)), max=float(h[m].max()),
                                                       largest_eta_H=float(eta[m].max())))
    order = np.argsort(-h)
    out['above_electron_gas_value'] = table([i for i in order if h[i] > GAS], n=100)

    # one proton in an electron gas at each compound's own valence-electron density
    gas = sorted((r['rs'], r['h_eV_A']) for r in json.load(open(ROOT / 'numerics' / 'jellium.json'))['records'] if r['Z'] == 1 and r['xc'] == 'pz81')
    desc = {r['id']: r for r in csv.DictReader(open(ROOT / 'model' / 'data' / 'descriptors.csv'))}
    rs = np.array([float(desc[r['id']]['r_s']) for r in lab])
    h_gas = np.interp(rs, [g[0] for g in gas], [g[1] for g in gas])
    inside = (rs >= gas[0][0]) & (rs <= gas[-1][0])
    lr = np.log(h[inside] / h_gas[inside])
    out['against_electron_gas_at_own_density'] = dict(
        note='r_s from all electrons outside the noble-gas core (three for lanthanides and actinides); a rule, not a calculation',
        n=int(inside.sum()), r_s_quantiles={str(q): float(np.percentile(rs, q)) for q in (5, 25, 50, 75, 95)},
        ratio_h_over_gas=dict(geometric_mean=float(np.exp(lr.mean())), spread_factor=float(np.exp(lr.std())),
                              quantiles={str(q): float(np.exp(np.percentile(lr, q))) for q in (5, 25, 50, 75, 95, 99)},
                              fraction_above_one=float((lr > 0).mean()), count_above_one=int((lr > 0).sum())),
        correlation_ln_h_ln_gas=float(np.corrcoef(np.log(h[inside]), np.log(h_gas[inside]))[0, 1]),
        mean_abs_error_ln=dict(gas_at_own_density=float(np.abs(lr).mean()), gas_at_own_density_rescaled=float(np.abs(lr - lr.mean()).mean()),
                               constant_36_6=float(np.abs(np.log(h[inside] / GAS)).mean()),
                               fitted_constant=float(np.abs(np.log(h[inside]) - np.log(h[inside]).mean()).mean())))

    # counts by the release's own transition temperatures at mu* = 0.10, and the five largest eta_H at three smearings
    def above(col, pool):
        v = np.array([num(r[col]) for r in pool])
        return {str(t): int((v > t).sum()) for t in (39, 60, 77, 100, 120)}
    out['hydrides_above_Tc'] = dict(allen_dynes=dict(with_spectrum=above('tc_ad', hyd), labelled=above('tc_ad', lab)),
                                    eliashberg=dict(with_spectrum=above('tc_el', hyd), labelled=above('tc_el', lab)))
    out['labelled_prototypes'] = len({r['prototype'] for r in lab})
    top5 = np.argsort(-eta)[:5]
    out['largest_eta_H_by_smearing'] = [dict(formula=lab[i]['formula'], eta_H={'0.005': round(num(lab[i]['S_top_005']), 2), '0.030': round(float(eta[i]), 2), '0.050': round(num(lab[i]['S_top_050']), 2)},
                                             lam={'0.005': round(num(lab[i]['lam_005']), 2), '0.030': round(num(lab[i]['lam']), 2), '0.050': round(num(lab[i]['lam_050']), 2)},
                                             h_005=round(num(lab[i]['S_top_005']) / float(rho[i]), 0), h_050=round(num(lab[i]['S_top_050']) / float(rho[i]), 0),
                                             tc_allen_dynes=round(num(lab[i]['tc_ad']), 1), tc_eliashberg=round(num(lab[i]['tc_el']), 1)) for i in top5]
    cov = ROOT / 'model' / 'data' / 'coverage.json'
    if cov.exists():
        out['coverage'].update(json.load(open(cov)))

    # points for the figures
    json.dump(dict(rho=[round(float(x), 4) for x in rho], eta=[round(float(x), 4) for x in eta],
                   tc20=[int(x >= 20) for x in tc], formula=[r['formula'] for r in lab]),
              open(ROOT / 'numerics' / 'h_alexandria_points.json', 'w'), separators=(',', ':'))

    json.dump(out, open(ROOT / 'numerics' / 'h_alexandria.json', 'w'), indent=1)
    show = {k: out[k] for k in ('labelled_hydrides', 'eta_on_density', 'correlation_ln_h_ln_density')}
    show['h'] = {k: out['h'][k] for k in ('geometric_mean', 'spread_factor', 'max', 'count_above_electron_gas_value')}
    show['eta_H_max'] = out['eta_H']['max']
    show['census_matched'] = (out['census_of_draft_4']['matched'], out['census_of_draft_4']['median_abs_ln_ratio'])
    print(json.dumps(show, indent=1))


if __name__ == '__main__':
    main()