"""Time mepd's geodesic interpolation (chainhelpers.run_geodesic) on the 30
reactions in data/endpoints.xyz.gz and write each path, for the mepd checkout
on PYTHONPATH.

    PYTHONPATH=<mepd checkout> python reproduce/time_geodesic.py TAG OUTDIR

Writes OUTDIR/times_TAG.csv (idx, nimages, seconds) and
OUTDIR/paths_TAG.xyz.gz (every image, Bohr as mepd holds it, repr floats).
Calls run one after another on one thread, after one untimed warm-up call
(numba compiles its kernel on first use).
"""
import csv, gzip, io, os, sys, time
from pathlib import Path

os.environ.setdefault("OMP_NUM_THREADS", "1")
import numpy as np
from qcconst.constants import ANGSTROM_TO_BOHR
from qcdata import Structure

from mepd import chainhelpers as ch
from mepd.nodes.node import StructureNode

HERE = Path(__file__).resolve().parent.parent
NIMAGES = (12, 30)


def read_pairs():
    lines = gzip.open(HERE / "data" / "endpoints.xyz.gz", "rt").read().splitlines()
    frames, i = [], 0
    while i < len(lines):
        n = int(lines[i]); head = lines[i + 1].split()
        atoms = [l.split() for l in lines[i + 2:i + 2 + n]]
        q = {k: int(v) for k, v in (t.split("=") for t in head[3:])}
        frames.append((head[0], head[2], [a[0] for a in atoms],
                       np.array([[float(x) for x in a[1:4]] for a in atoms]), q))
        i += 2 + n
    pairs = {}
    for idx, role, sym, xyz, q in frames:
        node = StructureNode(structure=Structure(symbols=sym, geometry=xyz * ANGSTROM_TO_BOHR,
                                                 charge=q["charge"], multiplicity=q["mult"]))
        pairs.setdefault(idx, {})[role] = node
    return [(idx, p["start"], p["end"]) for idx, p in pairs.items()]


def main():
    tag, out = sys.argv[1], Path(sys.argv[2])
    out.mkdir(parents=True, exist_ok=True)
    pairs = read_pairs()
    ch.run_geodesic([pairs[0][1], pairs[0][2]], nimages=12)       # warm-up
    times = []
    with io.TextIOWrapper(gzip.GzipFile(out / f"paths_{tag}.xyz.gz", "wb", mtime=0), encoding="utf-8") as fh:
        for idx, a, b in pairs:
            for nim in NIMAGES:
                t0 = time.perf_counter()
                chain = ch.run_geodesic([a, b], nimages=nim)
                times.append((idx, nim, time.perf_counter() - t0))
                for k, n in enumerate(chain.nodes):
                    fh.write(f"{len(n.coords)}\n{idx} nimages={nim} image={k}\n")
                    for s, x in zip(n.structure.symbols, np.asarray(n.coords)):
                        fh.write(f"{s} {float(x[0])!r} {float(x[1])!r} {float(x[2])!r}\n")
    with open(out / f"times_{tag}.csv", "w", newline="") as fh:
        w = csv.writer(fh); w.writerow(["idx", "nimages", "seconds"]); w.writerows(times)
    print(tag, f"total {sum(t for *_, t in times):.2f} s")


if __name__ == "__main__":
    main()