"""Rerun both versions once on all 60 calls and compare with the stored data.

    MEPD_BEFORE=<checkout at 1b8d63f> MEPD_AFTER=<checkout at b366826> \
        python reproduce/quick_check.py [--out result.json]

Passes when every path from both versions matches data/paths.xyz.gz within
1e-8 Bohr per coordinate (bit-identical on the original machine) and the total
speed-up is at least 1.5x. Timings depend on the machine; the stored ones are in
data/timings.csv. About a minute.
"""
import gzip, json, os, subprocess, sys, tempfile
from pathlib import Path

import numpy as np

HERE = Path(__file__).resolve().parent.parent


def read_paths(fn):
    lines = gzip.open(fn, "rt").read().splitlines()
    out, i = {}, 0
    while i < len(lines):
        n = int(lines[i]); head = lines[i + 1].split()
        xyz = [[float(v) for v in l.split()[1:4]] for l in lines[i + 2:i + 2 + n]]
        out.setdefault((head[0], head[1]), []).append(xyz)
        i += 2 + n
    return {k: np.array(v) for k, v in out.items()}


def main():
    out_json = sys.argv[sys.argv.index("--out") + 1] if "--out" in sys.argv else None
    stored = read_paths(HERE / "data" / "paths.xyz.gz")
    res, tmp = {}, Path(tempfile.mkdtemp())
    for tag in ("before", "after"):
        env = dict(os.environ, PYTHONPATH=os.environ[f"MEPD_{tag.upper()}"], OMP_NUM_THREADS="1")
        subprocess.run([sys.executable, str(HERE / "reproduce" / "time_geodesic.py"), tag, str(tmp)], env=env, check=True)
        got = read_paths(tmp / f"paths_{tag}.xyz.gz")
        diff = max(float(np.abs(got[k] - stored[k]).max()) for k in stored)
        total = sum(float(l.split(",")[2]) for l in (tmp / f"times_{tag}.csv").read_text().splitlines()[1:])
        res[tag] = {"max_coord_diff_bohr": diff, "bit_identical": diff == 0.0, "total_s": round(total, 2)}
    res["total_speedup"] = round(res["before"]["total_s"] / res["after"]["total_s"], 2)
    res["pass"] = all(res[t]["max_coord_diff_bohr"] <= 1e-8 for t in ("before", "after")) and res["total_speedup"] >= 1.5
    print(json.dumps(res, indent=1))
    if out_json:
        Path(out_json).write_text(json.dumps(res, indent=1) + "\n")
    sys.exit(0 if res["pass"] else 1)


if __name__ == "__main__":
    main()