"""wB97X-D/def2-TZVP single points with Psi4.

    python dft_single_points.py --out ts_sp.json ts_34.xyz ts_58.xyz ts_61.xyz
    python dft_single_points.py --split --out r_sp.json reactants.xyz

--split: each molecule of the structure separately (bonds from covalent radii), energies summed.
Writes {file name: energy in Hartree}. Closed shell, neutral.
"""
import argparse
import json
from pathlib import Path

import numpy as np
import psi4

COV = {"H": .31, "C": .76, "N": .71, "O": .66}
ELECTRONS = {"H": 1, "C": 6, "N": 7, "O": 8}

ap = argparse.ArgumentParser()
ap.add_argument("files", nargs="+")
ap.add_argument("--split", action="store_true")
ap.add_argument("--out", required=True)
ap.add_argument("--threads", type=int, default=4)
a = ap.parse_args()
psi4.set_num_threads(a.threads)
psi4.set_memory("4 GB")
psi4.core.be_quiet()


def read(fp):
    lines = Path(fp).read_text().split("\n")
    n = int(lines[0])
    rows = [line.split() for line in lines[2:2 + n]]
    return [r[0] for r in rows], np.array([[float(v) for v in r[1:4]] for r in rows])


def molecules(sym, xyz):
    parent = list(range(len(sym)))

    def root(i):
        while parent[i] != i:
            i = parent[i]
        return i

    for i in range(len(sym)):
        for j in range(i):
            if np.linalg.norm(xyz[i] - xyz[j]) < 1.25 * (COV[sym[i]] + COV[sym[j]]):
                parent[root(i)] = root(j)
    groups = {}
    for i in range(len(sym)):
        groups.setdefault(root(i), []).append(i)
    return list(groups.values())


def energy(sym, xyz):
    assert sum(ELECTRONS[s] for s in sym) % 2 == 0, "closed-shell molecules only"
    mol = psi4.geometry("0 1\n" + "\n".join(f"{s} {x:.8f} {y:.8f} {z:.8f}" for s, (x, y, z) in zip(sym, xyz))
                        + "\nsymmetry c1\nno_com\nno_reorient")
    psi4.set_options({"reference": "rks"})
    return psi4.energy("wb97x-d/def2-tzvp", molecule=mol)


out = {}
for fp in a.files:
    sym, xyz = read(fp)
    parts = molecules(sym, xyz) if a.split else [list(range(len(sym)))]
    out[Path(fp).name] = sum(energy([sym[k] for k in p], xyz[p]) for p in parts)
Path(a.out).write_text(json.dumps(out) + "\n")
print(json.dumps(out))