"""Origin-independent structure types, and one record per compound.

model/dataset.py labels a structure by the Wyckoff letters spglib assigns, and those
depend on which of several equivalent origins a record is described from, so one
structure type can carry two labels (cubic A2MH6 appears as 225_1a_2c_6e and as
225_1b_2c_6e). Here the label is built from the site-symmetry symbol of each occupied
orbit in place of its letter. Positions that a change of origin maps onto each other
have the same site symmetry, so the label no longer depends on the origin. It is
coarser than the Wyckoff label where two inequivalent positions share a site symmetry,
which makes a hold-out by structure type stricter.

Output: model/data/structure_types.csv with, for every compound of the training table,
  structure_type   space group, and for each element its count per formula unit and the
                   site symmetries it occupies, element names removed
  compound         reduced formula and structure type; records sharing it are the same compound
  keep             1 for the first record of a compound, 0 for later records of it

Run:  .venv/bin/python model/structure_types.py
"""
import csv
import pathlib
import pickle
from collections import Counter
from functools import reduce
from math import gcd

import spglib
from pymatgen.core import Composition, Element

ROOT = pathlib.Path(__file__).resolve().parents[1]


def structure_type(lattice, frac, species):
    numbers = [Element(s).Z for s in species]
    ds = spglib.get_symmetry_dataset((lattice, frac, numbers), symprec=0.05)
    if ds is None:
        return 'unresolved'
    counts = Counter(species)
    g = reduce(gcd, counts.values())
    per = {}
    for s, sym, orbit in zip(species, ds.site_symmetry_symbols, ds.crystallographic_orbits):
        per.setdefault(s, {})[int(orbit)] = sym.replace(' ', '')
    parts = sorted('%d[%s]' % (counts[s] // g, ','.join(sorted(per[s].values()))) for s in counts)
    return '%d_%s' % (ds.number, '_'.join(parts))


def main():
    rows = pickle.load(open(ROOT / 'data' / 'alexandria_ph' / 'dataset.pkl', 'rb'))
    rows = [r for r in rows if r['row']['S'] > 0]
    seen, out = set(), []
    for r in rows:
        st = structure_type(r['lattice'], r['frac'], r['species'])
        comp = Composition(r['row']['formula']).reduced_formula + '|' + st
        out.append((r['row']['id'], st, comp, int(comp not in seen)))
        seen.add(comp)
    with open(ROOT / 'model' / 'data' / 'structure_types.csv', 'w', newline='') as f:
        w = csv.writer(f)
        w.writerow(['id', 'structure_type', 'compound', 'keep'])
        w.writerows(out)
    lab = {r['row']['id'] for r in rows if r['row']['branch_ok'] and r['row']['h'] and r['row']['h'] > 0}
    kept = [o for o in out if o[3]]
    print('compounds', len(out), 'kept', len(kept), 'structure types', len({o[1] for o in out}),
          '| labelled hydrides', len(lab), 'kept', sum(1 for o in kept if o[0] in lab),
          'structure types among them', len({o[1] for o in kept if o[0] in lab}))


if __name__ == '__main__':
    main()