"""Build screen.csv (the table in data/) from a screened workspace.

    python screen_table.py ws [OUT.csv]

One row per screened reaction. `kept` is benchmark.py's keep rule: status match, and minimizing changed
neither IRC end into another species. `irc_retried` is 1 when the first IRC (fmax 0.1 eV/A)
had an end that changed species and the IRC was retried at fmax 0.01 (`benchmark.py reirc`), which then decides
the row ("failed" when the retry did not finish, so the first IRC stands); `loose_irc_status` is the first IRC's status. Needs only Python 3.
"""
import csv
import json
import sys
from pathlib import Path

WS = Path(sys.argv[1])
OUT = Path(sys.argv[2] if len(sys.argv) > 2 else "screen.csv")

rx = {r["id"]: r for r in csv.DictReader(open(WS / "reactions.csv"))}
rows = []
for ln in open(WS / "screen" / "results.jsonl"):
    if not ln.strip():
        continue
    r = json.loads(ln)
    ends = r.get("irc_ends") or []
    if r.get("irc_reversed"):
        ends = ends[::-1]
    connects = all(e.get("same_species_as_irc_end", True) for e in ends)
    drop = lambda k: round(ends[k]["drop_on_minimization_kcal"], 2) if len(ends) == 2 else ""
    freqs = r.get("ts_lowest_freqs_cm") or []
    x = rx[r["id"]]
    rows.append({
        "reaction_id": r["id"], "status": r["status"], "kept": int(r["status"] == "match" and connects),
        "irc_ends_kept_species": int(connects) if ends else "",
        "n_atoms": x["n_atoms"], "n_heavy": x["n_heavy"],
        "ts_imaginary_cm": round(freqs[0], 1) if freqs else "",
        "start_end_drop_kcal": drop(0), "end_end_drop_kcal": drop(1),
        "barrier_kcal": "" if r.get("barrier_kcal") is None else round(r["barrier_kcal"], 2),
        "dataset_barrier_kcal": round(float(x["reference_barrier_kcal"]), 2),
        "seconds": r.get("seconds", ""),
        "irc_retried": ("failed" if "status" in r["irc_retry"] else 1) if "irc_retry" in r else 0,
        "loose_irc_status": r.get("irc_loose", {}).get("status", ""),
    })
rows.sort(key=lambda r: r["reaction_id"])
with open(OUT, "w", newline="") as fh:
    w = csv.DictWriter(fh, fieldnames=list(rows[0]))
    w.writeheader()
    w.writerows(rows)
print(len(rows), "rows,", sum(r["kept"] for r in rows), "kept")