"""Recompute the g-xTB energy of every stored TS and compare the gaps with data/scores.csv.

    python reproduce/quick_check.py [--jobs 3] [--out result.json]     (needs g-xTB; set GXTB=/path/to/xtb if it is not on PATH as gxtb-xtb)

For each reaction with a lower TS, data/lower_ts.xyz.gz holds the reference g-xTB TS
and the TS(s) of the best route found. The gap is E(reference) minus the highest TS of
the route; it must match `lower_by_kcal` within 0.05 kcal/mol. This checks the numbers,
not that the structures are saddles (that was verified by TS optimization and IRC when
the sweeps ran).
"""

from __future__ import annotations

import argparse
import collections
import csv
import gzip
import json
import os
import subprocess
import tempfile
from concurrent.futures import ProcessPoolExecutor
from pathlib import Path

HERE = Path(__file__).resolve().parent
DATA = HERE.parent / "data"
HARTREE_TO_KCAL = 627.509474
TOL_KCAL = 0.05
# The xtb program from the g-xTB release (https://github.com/grimme-lab/g-xtb), version 2.0.1.
GXTB = os.environ.get("GXTB", "gxtb-xtb")


def frames():
    lines = gzip.open(DATA / "lower_ts.xyz.gz", "rt").read().splitlines()
    i = 0
    while i < len(lines):
        n = int(lines[i])
        rid, role = lines[i + 1].split()[:2]
        yield rid, role, lines[i + 2: i + 2 + n]
        i += 2 + n


def energy(args):
    """g-xTB total energy (Eh), run the way mepd runs it: one single-threaded gradient call."""
    atoms, charge, mult = args
    with tempfile.TemporaryDirectory() as d:
        Path(d, "ts.xyz").write_text(f"{len(atoms)}\n\n" + "\n".join(atoms) + "\n")
        cmd = [GXTB, "ts.xyz", "--silent", "--chrg", str(charge), "--grad"]
        if mult > 1:
            cmd += ["--uhf", str(mult - 1)]
        if Path(GXTB).name != "gxtb":
            cmd.append("--gxtb")
        subprocess.run(cmd, cwd=d, check=True, capture_output=True, env={**os.environ, "OMP_NUM_THREADS": "1"})
        for line in Path(d, "energy").read_text().splitlines():
            f = line.split()
            if len(f) >= 2 and f[0].isdigit():
                return float(f[1])
    raise RuntimeError("no energy in g-xTB output")


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--jobs", type=int, default=3)
    ap.add_argument("--out", type=Path)
    a = ap.parse_args()
    rows = {r["reaction_id"]: r for r in csv.DictReader(open(DATA / "scores.csv", newline=""))}
    fr = list(frames())
    with ProcessPoolExecutor(a.jobs) as pool:
        es = list(pool.map(energy, [(at, int(rows[rid]["charge"]), int(rows[rid]["multiplicity"])) for rid, _, at in fr]))
    by = collections.defaultdict(dict)
    for (rid, role, _), e in zip(fr, es):
        by[rid].setdefault("found", [])
        if role == "reference":
            by[rid]["reference"] = e
        else:
            by[rid]["found"].append(e)
    out, bad = [], 0
    for rid, d in by.items():
        gap = (d["reference"] - max(d["found"])) * HARTREE_TO_KCAL
        stored = float(rows[rid]["lower_by_kcal"])
        ok = abs(gap - stored) <= TOL_KCAL
        bad += not ok
        out.append({"reaction_id": rid, "recomputed_kcal": round(gap, 3), "stored_kcal": stored, "ok": ok})
    summary = {"n_reactions": len(out), "n_ok": len(out) - bad, "tolerance_kcal": TOL_KCAL,
               "max_abs_diff_kcal": round(max(abs(o["recomputed_kcal"] - o["stored_kcal"]) for o in out), 4)}
    print(json.dumps(summary))
    if a.out:
        a.out.write_text(json.dumps({"summary": summary, "reactions": out}, indent=1))
    return 0 if not bad else 1


if __name__ == "__main__":
    raise SystemExit(main())