"""Reads the Alexandria electron-phonon release and writes one compact record per compound.

Input:  data/alexandria_ph/alexandria_ph_NNN.json.bz2 (fetched by model/fetch_alexandria.sh)
Output: data/alexandria_ph/extract/NNN.jsonl.gz, one JSON object per line

Every record keeps the cell, the hydrogen density and a summary of the phonon frequencies.
Records with a spectral function also keep, at each of the ten smearing widths of the
release, the coupling constant, the logarithmic and second-moment frequencies and the
Hopfield sum

    S = M_H * lambda * <omega^2> = 2 M_H * integral of alpha^2F(omega) * omega d omega,

in eV/A^2 (the hydrogen mass is the unit of mass, as in the paper), and the same sum
restricted to the 3 n_H highest branches. At the width of 0.030 Ry, which is the one the
release uses for its own transition temperatures, the record also keeps lambda and S of
each branch.

Run:  .venv/bin/python model/extract.py [workers]
"""
import bz2
import gzip
import json
import pathlib
import sys
from concurrent.futures import ProcessPoolExecutor

import numpy as np

ROOT = pathlib.Path(__file__).resolve().parents[1]
SRC = ROOT / 'data' / 'alexandria_ph'
OUT = SRC / 'extract'

RY_EV = 13.605693123
HBAR2_OVER_MH = 4.14720e-3   # eV A^2: hbar^2 / M_H, so M_H omega^2 = (hbar omega)^2 / this
WIDTHS = [0.005 * i for i in range(1, 11)]   # Ry; 'a2F.dosN' is width N * 0.005
REF = 6                                      # index of the 0.030 Ry width


def moments(w, f):
    """lambda, omega_log (eV), omega_2 (eV) and S (eV/A^2) of one spectral function."""
    m = w > 1e-5
    w, f = w[m], f[m]
    lam = 2 * np.trapezoid(f / w, w)
    first = 2 * np.trapezoid(f * w, w)          # lambda <omega^2>, Ry^2
    if not (lam > 0 and first > 0):
        return None
    wlog = float(np.exp(2 / lam * np.trapezoid(f / w * np.log(w), w)))
    return dict(lam=float(lam), wlog=wlog * RY_EV, w2=float(np.sqrt(first / lam)) * RY_EV,
                S=float(first) * RY_EV ** 2 / HBAR2_OVER_MH)


def hydrogen_weights(dyns, is_h):
    """Mass-weighted share of each branch's eigenvector on hydrogen, per stored wavevector.

    The release stores displacement patterns normalised to one; multiplying by the square
    root of the mass and renormalising gives the eigenvector of the dynamical matrix.
    """
    mass = dict(dyns['species'])
    m = np.array([mass[s[0]] for s in dyns['sites']])
    out, freqs = [], []
    for q in dyns['dyn']:
        v = np.array(q['vecsR']) + 1j * np.array(q['vecsIM'])       # branch, atom, xyz
        p = (np.abs(v) ** 2).sum(axis=2) * m[None, :]
        out.append(p[:, is_h].sum(axis=1) / p.sum(axis=1))
        freqs.append(q['freqs'])
    return np.array(out), np.array(freqs)                            # q, branch


def record(e):
    d, st = e['data'], e['structure']
    species = [s['species'][0]['element'] for s in st['sites']]
    n_h = species.count('H')
    vol = st['lattice']['volume']
    r = dict(id=d['mat_id'], formula=d['formula'], nsites=d['nsites'],
             lattice=[[round(x, 6) for x in row] for row in st['lattice']['matrix']],
             species=species, frac=[[round(x, 6) for x in s['abc']] for s in st['sites']],
             volume=round(vol, 5), n_H=n_h, rho_H=n_h / vol, imag=bool(d['imag']),
             fermi=d.get('fermi_level'), k_coarse=d.get('kpoints_coarse'), k_fine=d.get('kpoints_fine'))

    dyns = d.get('dyns')
    if not isinstance(dyns, dict):               # a few records carry the string 'Failed' here
        r['dyn_error'] = str(dyns)[:40] if dyns else 'missing'
        dyns = None
    if dyns and dyns.get('dyn'):
        order = [s[0] for s in dyns['sites']]
        is_h = np.array([s == 'H' for s in order])
        try:
            hw, fr = hydrogen_weights(dyns, is_h)
            r.update(nq=len(fr), qgrid=dyns.get('qgrid'), w_min=float(fr.min()), w_max=float(fr.max()))
            if n_h:
                top = hw[:, -3 * n_h:]                 # the 3 n_H highest branches at each wavevector
                r.update(h_top_min=float(top.min()), h_top_mean=float(top.mean()),
                         h_rest_max=float(hw[:, :-3 * n_h].max()) if 3 * n_h < hw.shape[1] else 0.0)
            if d.get('a2f'):
                r['freqs_q'] = [[round(float(x), 2) for x in row] for row in fr]
                r['h_weight_branch'] = [round(float(x), 4) for x in hw.mean(axis=0)]
        except Exception as err:                        # malformed block: keep the structure
            r['dyn_error'] = str(err)[:80]

    a2f = d.get('a2f')
    if isinstance(a2f, dict) and a2f:
        widths = {}
        for i, width in enumerate(WIDTHS, start=1):
            blk = a2f.get(f'a2F.dos{i}')
            if not blk:
                continue
            w = np.array(blk['freqs'], dtype=float)
            mo = moments(w, np.array(blk['total'], dtype=float))
            if mo is None:
                continue
            modes = np.array(blk['modes'], dtype=float)           # frequency, branch
            m = w > 1e-5
            s_branch = 2 * np.trapezoid(modes[m] * w[m][:, None], w[m], axis=0) * RY_EV ** 2 / HBAR2_OVER_MH
            if n_h:
                mo['S_top'] = float(s_branch[-3 * n_h:].sum())
            b = (d.get('broadening') or {}).get(f'Broadening_{width:.4f}') or {}
            mo['dos'] = b.get('dos(Ef)')
            mo['lam_release'] = b.get('lambda')
            widths[f'{width:.3f}'] = {k: (round(v, 6) if isinstance(v, float) else v) for k, v in mo.items()}
            if i == REF:
                lam_branch = 2 * np.trapezoid(modes[m] / w[m][:, None], w[m], axis=0)
                r['S_branch'] = [round(float(x), 5) for x in s_branch]
                r['lam_branch'] = [round(float(x), 5) for x in lam_branch]
        r['a2f'] = widths
        tc = d.get('tc') or {}
        mu = tc.get('mustr') or []
        pick = lambda key: {f'{m:.2f}': v for m, v in zip(mu, tc.get(key) or []) if 0.09 < m < 0.17}
        r['tc'] = dict(lam=tc.get('lambda'), wlog_K=tc.get('wlog[K]'), w2_K=tc.get('w2av[K]'),
                       allen_dynes=pick('TcAllenDynes'), eliashberg=pick('TcEliashberg'))
    return r


def one(path):
    out = OUT / (path.name.split('_')[-1].split('.')[0] + '.jsonl.gz')
    if out.exists():
        return path.name, None
    with bz2.open(path, 'rt') as f:
        entries = json.load(f)['entries']
    rows = [record(e) for e in entries]
    tmp = out.with_suffix('.tmp')
    with gzip.open(tmp, 'wt') as g:
        for r in rows:
            g.write(json.dumps(r, separators=(',', ':')) + '\n')
    tmp.rename(out)
    return path.name, (len(rows), sum(1 for r in rows if 'a2f' in r), sum(1 for r in rows if r['n_H']),
                       sum(1 for r in rows if r['n_H'] and 'a2f' in r))


if __name__ == '__main__':
    OUT.mkdir(exist_ok=True)
    files = sorted(SRC.glob('alexandria_ph_*.json.bz2'))
    if len(sys.argv) > 2:                                   # explicit file numbers
        files = [f for f in files if f.name.split('_')[-1].split('.')[0] in sys.argv[2:]]
    workers = int(sys.argv[1]) if len(sys.argv) > 1 else 3
    with ProcessPoolExecutor(workers) as pool:
        for name, stats in pool.map(one, files):
            print(name, 'skipped' if stats is None else 'records %d, with spectrum %d, hydrides %d, hydrides with spectrum %d' % stats, flush=True)