"""First use of roomtsc 0.1: estimates for the hydride records of the release that have no label.

Two thirds of the release's records, and 84% of its hydride records, have imaginary harmonic
modes and carry no spectral function. The first paper names them as its largest gap: strong
coupling drives phonons toward instability, so they may hold larger scattering strengths
than the labelled records. This script runs the released estimators on every such hydride
record and writes what they return.

What the output can and cannot show. The estimators were trained on dynamically stable
structures only, they failed the test of predicting above their training range, and an
unstable structure is not a minimum of the energy. The estimates rank these records for a
direct calculation of the Hopfield sum, which is defined whether or not a mode is unstable.
They are not evidence about what that calculation will find.

Input:  data/alexandria_ph/extract/*.jsonl.gz, model/weights/
Output: model/results2/unlabelled_hydrides.csv   one row per compound
        model/results2/unlabelled_hydrides.json  counts, distributions and the leading rows
Run:    .venv/bin/python model/screen_unlabelled.py [workers]
"""
import csv
import glob
import gzip
import json
import math
import pathlib
import sys
from concurrent.futures import ProcessPoolExecutor

import joblib
import numpy as np
import torch
from pymatgen.core import Composition

HERE = pathlib.Path(__file__).resolve().parent
ROOT = HERE.parent
sys.path.insert(0, str(HERE))
import dataset as D
import structure_types as ST
import train as T

W = HERE / 'weights'
META = json.load(open(W / 'meta.json'))
ORDER = json.load(open(W / 'feature_order.json'))
SEEN = set(META['elements_seen'])


def prep(r):
    """Graph and descriptors of one record, or the reason it is skipped."""
    species = r['species']
    if set(species) - SEEN:
        return dict(skip='element outside the training set')
    try:
        i, j, d = D.graph(r['lattice'], r['frac'], species)
        sg, _ = D.prototype(r['lattice'], r['frac'], species)
        stype = ST.structure_type(r['lattice'], r['frac'], species)
        feats = D.describe(r)
        feats.update(D.geometry(species, i, j, d))
        feats['spacegroup'] = sg
    except Exception as err:
        return dict(skip='could not build the graph: ' + str(err)[:60])
    ev = [T.element_vector(s) for s in species]
    g = dict(x=np.stack([e[0] for e in ev]), z=np.array([e[1] for e in ev]),
             inv_m=np.array([T.M_H / e[2] for e in ev], dtype=np.float32),
             is_h=np.array([s == 'H' for s in species]), i=i, j=j, d=d,
             vol=r['volume'], ln_S=0.0, ln_h=float('nan'))
    return dict(id=r['id'], formula=r['formula'], n_sites=len(species), rho_H=r['rho_H'], imag=r['imag'],
                w_min=r.get('w_min'), stype=stype, compound=Composition(r['formula']).reduced_formula + '|' + stype,
                x=[feats[k] for k in ORDER], g=g)


def load(path):
    """Hydride records of one file that carry no spectral function."""
    out, counts = [], dict(hydride_records=0, with_spectrum=0)
    with gzip.open(path, 'rt') as f:
        for line in f:
            r = json.loads(line)
            if not r['n_H']:
                continue
            counts['hydride_records'] += 1
            if 'a2f' in r:
                counts['with_spectrum'] += 1
                continue
            out.append(prep(r))
    return out, counts


@torch.no_grad()
def main():
    workers = int(sys.argv[1]) if len(sys.argv) > 1 else 8
    files = sorted(glob.glob(str(ROOT / 'data' / 'alexandria_ph' / 'extract' / '*.jsonl.gz')))
    rows, counts, skipped = [], dict(hydride_records=0, with_spectrum=0), {}
    with ProcessPoolExecutor(workers) as pool:
        for part, c in pool.map(load, files):
            for k in counts:
                counts[k] += c[k]
            for p in part:
                if 'skip' in p:
                    skipped[p['skip'].split(':')[0]] = skipped.get(p['skip'].split(':')[0], 0) + 1
                else:
                    rows.append(p)
    index = json.load(open(W / 'training_index.json'))
    seen, keep, in_training = set(), [], 0
    for r in rows:                                   # one row per compound, and none that training holds
        if r['compound'] in index:
            in_training += 1
            continue
        if r['compound'] in seen:
            continue
        seen.add(r['compound'])
        keep.append(r)
    rows = keep

    nets = []
    for seed in META['graph_seeds']:
        net = T.Net()
        net.load_state_dict(torch.load(W / f'graph_{seed}.pt', map_location='cpu'))
        net.eval()
        nets.append(net)
    trees = joblib.load(W / 'trees.joblib')
    lt = trees.predict(np.array([r['x'] for r in rows], dtype=float))
    lg = np.zeros(len(rows))
    for k in range(0, len(rows), 256):
        batch = T.collate(rows[k:k + 256], 'cpu')
        lg[k:k + 256] = np.mean([net(batch)[1].numpy() for net in nets], axis=0)
    lh = 0.5 * (lt + lg)
    h = np.exp(lh)
    rho = np.array([r['rho_H'] for r in rows])
    eta = rho * h
    types = set(META['structure_types_seen_in_labelled_hydrides'])
    known = np.array([r['stype'] in types for r in rows])
    wmin = np.array([r['w_min'] if r['w_min'] is not None else np.nan for r in rows])

    out = ROOT / 'model' / 'results2'
    with open(out / 'unlabelled_hydrides.csv', 'w', newline='') as f:
        w = csv.writer(f)
        w.writerow(['id', 'formula', 'n_sites', 'rho_H', 'lowest_frequency_cm', 'h_trees', 'h_graph', 'h', 'eta_H', 'structure_type', 'structure_type_in_training'])
        for i in np.argsort(-eta):
            r = rows[i]
            w.writerow([r['id'], r['formula'], r['n_sites'], '%.4f' % rho[i], '' if np.isnan(wmin[i]) else '%.0f' % wmin[i],
                        '%.2f' % math.exp(lt[i]), '%.2f' % math.exp(lg[i]), '%.2f' % h[i], '%.3f' % eta[i], r['stype'], int(known[i])])

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

    # the same estimate for the labelled hydrides, out of sample: the grouped folds of the second run
    g = json.load(open(out / 'grouped.json'))['compounds_held_out']
    est_lab = np.exp(0.5 * (np.log([c['trees'] for c in g]) + np.log([c['graph'] for c in g])))
    lab = np.array([c['h'] for c in g])

    def table(order, n=25):
        return [dict(formula=rows[i]['formula'], id=rows[i]['id'], rho_H=round(float(rho[i]), 4), h=round(float(h[i]), 1),
                     h_trees=round(math.exp(lt[i]), 1), h_graph=round(math.exp(lg[i]), 1), eta_H=round(float(eta[i]), 2),
                     lowest_frequency_cm=None if np.isnan(wmin[i]) else round(float(wmin[i])),
                     structure_type_in_training=bool(known[i])) for i in order[:n]]

    # the part of the list the estimators can speak to: a second element, and a hydrogen
    # density inside the range they were trained on
    rho_max = META['training_range']['rho_H'][1]
    pure = np.array([set(r['g']['z'].tolist()) == {1} for r in rows])
    inside = ~pure & (rho <= rho_max)
    order_inside = [i for i in np.argsort(-eta) if inside[i]]
    restricted = dict(note='compounds with an element besides hydrogen and a hydrogen density no higher than the largest in training (%.3f per A^3)' % rho_max,
                      n=int(inside.sum()), hydrogen_only=int(pure.sum()), above_training_density=int((~pure & (rho > rho_max)).sum()),
                      estimate_h=geo(h[inside]), estimate_eta_H=dict(geo(eta[inside]), count_above={str(t): int((eta[inside] > t).sum()) for t in (1, 2, 3, 4.83)}),
                      estimate_h_count_above={str(t): int((h[inside] > t).sum()) for t in (20, 30, 36.6)},
                      largest_eta_H=table(order_inside, 20))
    ok = ~np.isnan(wmin)
    summary = dict(
        model='roomtsc ' + META['version'], counts=dict(counts, without_spectrum=counts['hydride_records'] - counts['with_spectrum'],
                                                         skipped=skipped, same_compound_in_training=in_training, compounds_estimated=len(rows),
                                                         with_structure_type_in_training=int(known.sum())),
        estimate_h=geo(h), estimate_h_structure_type_in_training=geo(h[known]), estimate_h_structure_type_new=geo(h[~known]) if (~known).any() else None,
        labelled_out_of_sample=dict(estimate_h=geo(est_lab), label_h=geo(lab)),
        estimate_eta_H=dict(geo(eta), count_above={str(t): int((eta > t).sum()) for t in (1, 2, 3, 4.83)}),
        estimate_h_count_above={str(t): int((h > t).sum()) for t in (20, 30, 36.6)},
        rho_H=dict(geo(rho), count_above={str(t): int((rho > t).sum()) for t in (0.10, 0.12, 0.15, 0.20)}),
        instability=dict(note='rank correlation of the estimate of h with the depth of the lowest harmonic frequency (more negative is more unstable)',
                         spearman_h_vs_depth=float(T.spearman(np.log(h[ok]), -wmin[ok])), n=int(ok.sum()),
                         median_lowest_frequency_cm=float(np.median(wmin[ok]))),
        inside_training_range=restricted,
        largest_eta_H=table(np.argsort(-eta)), largest_h=table(np.argsort(-h)),
        largest_eta_H_known_structure_types=table([i for i in np.argsort(-eta) if known[i]]))
    json.dump(summary, open(out / 'unlabelled_hydrides.json', 'w'), indent=1)
    print(json.dumps({k: summary[k] for k in ('counts', 'estimate_h', 'labelled_out_of_sample', 'estimate_eta_H', 'estimate_h_count_above', 'instability')}, indent=1))


if __name__ == '__main__':
    main()