"""Rerun a random subset of the screen with today's mepd and compare per-reaction status.

    python reproduce/quick_rerun.py --n 48 --jobs 3

Needs mepd with g-xTB, and the Reaction-QM files B3LYPD3_TZVP.h5 and
B3LYPD3_TZVP_reaction_info.csv (doi:10.5281/zenodo.18551029) in
reproduce/rerun/data/reaction_qm/. The screen code is in reproduce/rerun/src/.
Writes provenance/reproduce_last.json (manifest plus
per-reaction comparison). Exit 0 when status agreement >= --min-agree.
"""

from __future__ import annotations

import argparse
import csv
import io
import zipfile
import json
import os
import random
import sys
import time
from pathlib import Path

HERE = Path(__file__).resolve().parent
FINDING = HERE.parent
RERUN = HERE / "rerun"


def main() -> int:
    p = argparse.ArgumentParser()
    p.add_argument("--n", type=int, default=48)
    p.add_argument("--jobs", type=int, default=3)
    p.add_argument("--seed", type=int, default=20261013)
    p.add_argument("--min-agree", type=float, default=0.95)
    a = p.parse_args()

    os.environ["OMP_NUM_THREADS"] = "1"
    sys.path.insert(0, str(RERUN / "src"))
    from benchmark import screen_gxtb as sg

    with zipfile.ZipFile(FINDING / "data" / "screen.csv.zip") as z, \
            io.TextIOWrapper(z.open("screen.csv"), encoding="utf-8", newline="") as fh:
        stored = {r["reaction_id"]: r for r in csv.DictReader(fh)}
    ids = random.Random(a.seed).sample(sorted(stored), a.n)

    import multiprocessing as mp

    started = time.time()
    ctx = mp.get_context("fork")
    with ctx.Pool(a.jobs, initializer=sg._init_worker, initargs=(str(RERUN / "mepd_profile.toml"), sg.DEFAULT_TIMEOUT_S)) as pool:
        new = {r["reaction_id"]: r for r in pool.imap_unordered(sg.screen_one, ids)}

    rows = []
    for rid in ids:
        old, cur = stored[rid], new[rid]
        d_old = old["gxtb_barrier_sep_kcal"]
        d_new = cur.get("gxtb_barrier_sep_kcal")
        rows.append({
            "reaction_id": rid, "stored": old["status"], "rerun": cur["status"],
            "barrier_stored": float(d_old) if d_old else None, "barrier_rerun": d_new,
        })
    agree = sum(r["stored"] == r["rerun"] for r in rows) / len(rows)
    both = [r for r in rows if r["barrier_stored"] is not None and r["barrier_rerun"] is not None]
    max_db = max((abs(r["barrier_stored"] - r["barrier_rerun"]) for r in both), default=None)

    result = {
        "label": "quick rerun", "n": a.n, "seed": a.seed, "jobs": a.jobs,
        "wall_s": round(time.time() - started, 1), "status_agreement": round(agree, 3),
        "max_abs_barrier_change_kcal": round(max_db, 3) if max_db is not None else None,
        "comparison": rows,
    }
    try:  # in the findings repo, also record versions and commits
        sys.path.insert(0, str(FINDING.parents[1] / "src"))
        from findings.manifest import host_manifest
        mepd_repo = Path(__import__("mepd").__file__).resolve().parents[1]
        out = host_manifest([mepd_repo], ["gxtb-xtb"], result)
    except ImportError:
        out = result
    (FINDING / "provenance").mkdir(exist_ok=True)
    (FINDING / "provenance" / "reproduce_last.json").write_text(json.dumps(out, indent=2))
    print(f"status agreement {agree:.1%} on {a.n}; max |Δbarrier| {max_db} kcal/mol; "
          f"{out['wall_s']} s")
    for r in rows:
        if r["stored"] != r["rerun"]:
            print(f"  {r['reaction_id']}: {r['stored']} -> {r['rerun']}")
    return 0 if agree >= a.min_agree else 1


if __name__ == "__main__":
    sys.exit(main())