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

    python screen_table.py ws [OUT.csv]

One row per screened reaction. Reads ws/reactions.csv, ws/screen/results.all-bonds-retried.jsonl (every bond
counts, metal-ligand contacts too, with the IRC retry decided the same way: screen_all_bonds.py) and
ws/screen/results.jsonl (the same calculations judged by covalent bonds only, after `benchmark.py recompare
--no-metal-ligand-contacts` and `benchmark.py reirc`). `kept` is benchmark.py's keep rule: status match, and
minimizing changed neither IRC end into another species. `irc_retried` is 1 when an end of the first IRC (fmax
0.1 eV/A) changed species under either comparison and the IRC was retried at fmax 0.01, 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 (all bonds). 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")
STRICT = WS / "screen" / "results.all-bonds-retried.jsonl"
COVALENT = WS / "screen" / "results.jsonl"


def load(p):
    out = {}
    for ln in open(p):
        if ln.strip():
            r = json.loads(ln)
            out[r["id"]] = r
    return out


def ends(r):
    e = r.get("irc_ends") or []
    return e[::-1] if r.get("irc_reversed") else e


def kept(r):
    return int(r.get("status") == "match" and all(e.get("same_species_as_irc_end", True) for e in ends(r)))


if __name__ == "__main__":
    rx = {r["id"]: r for r in csv.DictReader(open(WS / "reactions.csv"))}
    strict, cov = load(STRICT), load(COVALENT)
    rows = []
    for rid, r in strict.items():
        c, x, e = cov[rid], rx[rid], ends(r)
        connects = all(v.get("same_species_as_irc_end", True) for v in e)
        drop = lambda k: round(e[k]["drop_on_minimization_kcal"], 2) if len(e) == 2 else ""
        freqs = r.get("ts_lowest_freqs_cm") or []
        sides = c.get("sides") or {}
        mlc = [sides.get(s, {}).get("metal_ligand_changed") for s in ("start", "end")]
        rows.append({
            "reaction_id": rid, "metal": rid.rsplit("_", 1)[1], "multiplicity": x["multiplicity"],
            "status": r["status"], "kept": kept(r),
            "status_covalent": c["status"], "kept_covalent": kept(c),
            "metal_ligand_changed": int(any(mlc)) if any(v is not None for v in mlc) else "",
            "irc_ends_kept_species": int(connects) if e 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 (all bonds),", sum(r["kept_covalent"] for r in rows), "by covalent bonds only")