"""Quick check: recompute from data/structures.xyz the g-xTB energies of the separated reactants and of the
three transition states, and the endo/exo assignment of the three adducts, and compare them with
data/cycloaddition_ts.csv. Needs mepd (g-xTB) and RDKit; takes about a minute.

    python quick_rerun.py --data DATA_DIR [--out result.json]
"""
import argparse
import csv
import json
import sys
import tempfile
from pathlib import Path

import numpy as np

sys.path.insert(0, str(Path(__file__).parent))
from build_tables import BOHR_PER_ANGSTROM, stereo  # noqa: E402

ap = argparse.ArgumentParser()
ap.add_argument("--data", required=True, type=Path)
ap.add_argument("--out", type=Path)
a = ap.parse_args()

frames = {}
lines = (a.data / "structures.xyz").read_text().split("\n")
i = 0
while i < len(lines) and lines[i].strip():
    n = int(lines[i])
    frames[lines[i + 1].split(":")[0].split(" (")[0].strip()] = lines[i:i + 2 + n]
    i += 2 + n

import mepd  # noqa: E402
from qcdata import Structure  # noqa: E402

from mepd.discovery.free_energy import molecules  # noqa: E402
from mepd.engines.gxtb import GXTBCalculator  # noqa: E402
from mepd.nodes.node import StructureNode  # noqa: E402

eng = GXTBCalculator()


def parse(block):
    rows = [r.split() for r in block[2:]]
    return [r[0] for r in rows], np.array([[float(v) for v in r[1:4]] for r in rows])


def e(sym, x):
    node = StructureNode(structure=Structure(symbols=list(sym), geometry=np.asarray(x) * BOHR_PER_ANGSTROM))
    return float(mepd.energies([node], eng)[0])


sym, x = parse(frames["reactants"])
e_sep = sum(e([sym[k] for k in p], x[p]) for p in molecules(sym, x))
report, worst = [], 0.0
with open(a.data / "cycloaddition_ts.csv", newline="") as fh:
    for r in csv.DictReader(fh):
        k = r["adduct_species"]
        e_ts = e(*parse(frames[f"ts_{k}"]))
        with tempfile.NamedTemporaryFile("w", suffix=".xyz", delete=False) as tmp:
            tmp.write("\n".join(frames[f"species {k}"]) + "\n")
        st = stereo(tmp.name)
        d = max(abs(e_ts - float(r["gxtb_ts_hartree"])), abs(e_sep - float(r["gxtb_reactants_separated_hartree"])))
        worst = max(worst, d)
        report.append({"adduct": k, "stereo": st, "stereo_recorded": r["stereo"], "max_energy_diff_hartree": d})
ok = worst < 1e-6 and all(x["stereo"] == x["stereo_recorded"] for x in report)
result = {"ok": ok, "tolerance_hartree": 1e-6, "items": report}
print(json.dumps(result, indent=1))
if a.out:
    a.out.write_text(json.dumps(result, indent=1) + "\n")
sys.exit(0 if ok else 1)