"""Build F0001's tables (version 3: with the IRC retry of `benchmark.py reirc`) from the checked g-xTB screen (dataset-benchmark `screen`).

    python build_tables.py WS POPULATION_ZIP OUT_DIR

WS is the screened workspace (reactions.csv, screen/results.jsonl, screen/<id>/{start,end,ts}.xyz);
POPULATION_ZIP is version 1's data/population.csv.zip (all 199,890 B3LYP-RXN reactions; the same
63,954 are screened). Writes OUT_DIR/screen.csv.zip (one row per screened reaction) and
OUT_DIR/large/benchmark_set.xyz.gz (reactant, transition state and product of every kept reaction).
"""
import csv
import gzip
import io
import json
import sys
import zipfile
from pathlib import Path

ws, pop_zip, out = Path(sys.argv[1]), Path(sys.argv[2]), Path(sys.argv[3])
out.mkdir(parents=True, exist_ok=True)
(out / "large").mkdir(exist_ok=True)

with zipfile.ZipFile(pop_zip) as z:
    pop = {p["reaction_id"]: p for p in csv.DictReader(io.TextIOWrapper(z.open("population.csv"), encoding="utf-8"))}
rx = {r["id"]: r for r in csv.DictReader(open(ws / "reactions.csv", newline=""))}
screen = {}
for ln in open(ws / "screen" / "results.jsonl"):
    if ln.strip():
        r = json.loads(ln)
        screen[r["id"]] = r


def irc_connects(r):
    return all(e.get("same_species_as_irc_end", True) for e in r.get("irc_ends") or [])


FIELDS = ["reaction_id", "status", "kept", "n_atoms", "n_heavy", "elements", "n_reactant_frags", "balanced",
          "ts_imaginary_cm", "ts_second_lowest_cm", "irc_reversed",
          "start_end_drop_kcal", "end_end_drop_kcal", "start_end_changed_species", "end_end_changed_species",
          "start_matches", "end_matches", "start_reason", "end_reason", "undetermined_sides",
          "barrier_kcal", "reverse_barrier_kcal", "dataset_barrier_kcal", "seconds",
          "irc_start_smiles", "irc_end_smiles", "dataset_start_smiles", "dataset_end_smiles", "message",
          "irc_retried", "loose_irc_status", "irc_retry_seconds"]
rows = []
for rid in sorted(screen):
    r, p, x = screen[rid], pop[rid], rx[rid]
    fr = r.get("ts_lowest_freqs_cm") or []
    ends = list(r.get("irc_ends") or [])
    if len(ends) == 2 and r.get("irc_reversed"):        # in reaction order: reactant side first
        ends = ends[::-1]
    sides = r.get("sides") or {}
    st, en = sides.get("start", {}), sides.get("end", {})
    yes = lambda s: "" if not s else int(bool(s.get("graph") and s.get("stereo")))
    rnd = lambda v: "" if v is None else round(v, 3)
    rows.append({
        "reaction_id": rid, "status": r["status"], "kept": int(r["status"] == "match" and irc_connects(r)),
        "n_atoms": x["n_atoms"], "n_heavy": x["n_heavy"], "elements": p["elements"],
        "n_reactant_frags": p["n_reactant_frags"], "balanced": p["balanced"],
        "ts_imaginary_cm": fr[0] if fr else "", "ts_second_lowest_cm": fr[1] if len(fr) > 1 else "",
        "irc_reversed": "" if "irc_reversed" not in r else int(r["irc_reversed"]),
        "start_end_drop_kcal": rnd(ends[0]["drop_on_minimization_kcal"]) if ends else "",
        "end_end_drop_kcal": rnd(ends[1]["drop_on_minimization_kcal"]) if ends else "",
        "start_end_changed_species": int(not ends[0]["same_species_as_irc_end"]) if ends else "",
        "end_end_changed_species": int(not ends[1]["same_species_as_irc_end"]) if ends else "",
        "start_matches": yes(st), "end_matches": yes(en),
        "start_reason": st.get("reason", ""), "end_reason": en.get("reason", ""),
        "undetermined_sides": " ".join(r.get("undetermined_sides") or []),
        "barrier_kcal": rnd(r.get("barrier_kcal")), "reverse_barrier_kcal": rnd(r.get("reverse_barrier_kcal")),
        "dataset_barrier_kcal": x["reference_barrier_kcal"], "seconds": r.get("seconds", ""),
        "irc_start_smiles": (st.get("smiles") or ["", ""])[0], "dataset_start_smiles": (st.get("smiles") or ["", ""])[1],
        "irc_end_smiles": (en.get("smiles") or ["", ""])[0], "dataset_end_smiles": (en.get("smiles") or ["", ""])[1],
        "message": (r.get("message") or "").replace("\n", " ")[:300],
        "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", ""),
        "irc_retry_seconds": (r.get("irc_retry") or {}).get("seconds") or "",
    })

buf = io.StringIO()
w = csv.DictWriter(buf, fieldnames=FIELDS, lineterminator="\n")
w.writeheader()
w.writerows(rows)
with zipfile.ZipFile(out / "screen.csv.zip", "w", zipfile.ZIP_DEFLATED, compresslevel=9) as z:
    info = zipfile.ZipInfo("screen.csv", date_time=(1980, 1, 1, 0, 0, 0))   # byte-reproducible
    info.compress_type = zipfile.ZIP_DEFLATED
    z.writestr(info, buf.getvalue())

with io.TextIOWrapper(gzip.GzipFile(out / "large" / "benchmark_set.xyz.gz", "wb", compresslevel=9, mtime=0),
                      encoding="utf-8", newline="") as fh:
    for row in rows:
        if not row["kept"]:
            continue
        for part in ("start", "ts", "end"):
            lines = (ws / "screen" / row["reaction_id"] / f"{part}.xyz").read_text().splitlines()
            n = int(lines[0])
            fh.write(f"{n}\n{row['reaction_id']} {part}\n" + "\n".join(lines[2:2 + n]) + "\n")
print(f"{len(rows)} screened, {sum(r['kept'] for r in rows)} kept")