"""Screen a random sample of F0001's reactions again and compare with data/screen.csv.zip.

    python quick_rerun.py --h5 PATH/B3LYPD3_TZVP.h5 [--n 12] [--jobs 4] [--seed 20261022]

Imports the sample from the Reaction-QM file (doi:10.5281/zenodo.18551029) into a fresh workspace,
runs `benchmark.py screen` and `benchmark.py reirc` on it (uv, g-xTB 2.0.1 as `gxtb-xtb` on PATH), and compares each
reaction's outcome, whether it is kept, and its barrier with the table. Exits 1 if fewer than 90% of
outcomes agree.
"""
import argparse
import csv
import io
import json
import os
import random
import subprocess
import sys
import zipfile
from pathlib import Path

HERE = Path(__file__).resolve().parent
ap = argparse.ArgumentParser()
ap.add_argument("--h5", required=True, type=Path)
ap.add_argument("--n", type=int, default=12)
ap.add_argument("--jobs", type=int, default=4)
ap.add_argument("--seed", type=int, default=20261022)
ap.add_argument("--ws", type=Path, default=HERE / "quick_ws")
ap.add_argument("--also", default="", help="comma-separated reaction ids to screen as well (e.g. some with irc_retried 1)")
a = ap.parse_args()

with zipfile.ZipFile(HERE.parent.parent / "data" / "screen.csv.zip") as z:
    table = {r["reaction_id"]: r for r in csv.DictReader(io.TextIOWrapper(z.open("screen.csv"), encoding="utf-8"))}
ids = random.Random(a.seed).sample(sorted(table), a.n)
ids += [i for i in a.also.split(",") if i and i not in ids]
a.ws.mkdir(parents=True, exist_ok=True)
(a.ws / "ids.txt").write_text("\n".join(ids) + "\n")
env = dict(os.environ, OMP_NUM_THREADS="1")
uv = ["uv", "run", "--no-config", str(HERE / "benchmark.py")]
if not (a.ws / "reactions.csv").exists():
    subprocess.run(uv + ["import-reaction-qm", str(a.ws), str(a.h5), "--ids", str(a.ws / "ids.txt")], check=True, env=env)
subprocess.run(uv + ["screen", str(a.ws), "--jobs", str(a.jobs), "--timeout", "3600"], check=True, env=env)
subprocess.run(uv + ["reirc", str(a.ws), "--jobs", str(a.jobs), "--timeout", "1800"], check=True, env=env)

got = {}
for ln in open(a.ws / "screen" / "results.jsonl"):
    if ln.strip():
        r = json.loads(ln)
        got[r["id"]] = r
same = 0
for rid in ids:
    t, g = table[rid], got.get(rid, {})
    kept = int(g.get("status") == "match" and all(e.get("same_species_as_irc_end", True) for e in g.get("irc_ends") or []))
    b0 = float(t["barrier_kcal"]) if t["barrier_kcal"] else None
    b1 = g.get("barrier_kcal")
    db = "" if b0 is None or b1 is None else f"barrier {b1 - b0:+.4f} kcal/mol"
    ok = g.get("status") == t["status"] and str(kept) == t["kept"]
    same += ok
    print(f"{rid}  table {t['status']:<22} rerun {g.get('status', 'missing'):<22} {'same' if ok else 'DIFFERENT'}  {db}")
print(f"{same} of {len(ids)} outcomes agree")
sys.exit(0 if same >= 0.9 * len(ids) else 1)