"""Build search_outcomes.csv (the table in data/) from a searched workspace.

    python outcomes.py ws [OUT.csv]

One row per Transition1x-TMC reaction searched, in the columns of data/search_outcomes.csv. A reaction collected twice
keeps its real result over an empty error row. When `benchmark.py rescore` has run (ws/rescore/results.jsonl),
its best route, barrier, outcome and status replace the search's own (route steps whose IRC ends change
species are not counted). 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 "search_outcomes.csv")


def jl(p):
    return [json.loads(ln) for ln in p.read_text().splitlines() if ln.strip()]


rx = {r["id"]: r for r in csv.DictReader(open(ws / "reactions.csv"))}
screen = {r["id"]: r for r in jl(ws / "screen/results.jsonl")}
search = {}
for s in jl(ws / "search/results.jsonl"):
    if s["id"] not in search or s.get("status") != "error":
        search[s["id"]] = s
rescore = {r["id"]: r for r in jl(ws / "rescore/results.jsonl")} if (ws / "rescore/results.jsonl").exists() else {}
for i, r in rescore.items():
    if i in search:
        new = {k: r.get(k) for k in ("status", "outcome", "best_kind", "best_barrier_kcal", "lower_by_kcal")}
        search[i] = {**search[i], **new, "routes_dropped": r.get("routes_dropped", 0), "irc_retries": r.get("irc_retries", 0)}
rows = [{"reaction_id": s["id"], "metal": s["id"].rsplit("_", 1)[1], "n_atoms": rx[s["id"]]["n_atoms"], "status": s.get("status"),
         "outcome": s.get("outcome") if s.get("status") == "done" else "", "best_kind": s.get("best_kind") or "",
         "ref_barrier_kcal": s.get("ref_barrier_kcal"), "best_barrier_kcal": s.get("best_barrier_kcal"),
         "lower_by_kcal": s.get("lower_by_kcal"), "n_direct": s.get("n_direct"), "n_multistep": s.get("n_multistep"),
         "seconds": s.get("seconds"), "dataset_barrier_kcal": rx[s["id"]]["reference_barrier_kcal"],
         "screen_barrier_kcal": screen[s["id"]].get("barrier_kcal"), "routes_dropped": s.get("routes_dropped", ""),
         "step_ircs_retried": s.get("irc_retries", "")} for s in search.values()]
with open(out, "w", newline="") as fh:
    w = csv.DictWriter(fh, fieldnames=list(rows[0]))
    w.writeheader()
    w.writerows(rows)
print(f"{out}: {len(rows)} reactions")