"""Build search_diagnostics.csv (the table in data/) from a searched workspace: per reaction with a route
or none, whether the pair of the reference IRC's own two ends was searched, and what happened.

    python search_diagnostics.py ws search_outcomes.csv [OUT.csv]

Needs Python 3 with numpy.

Structures are compared by their sorted interatomic distances (same geometry up to atom numbering, 0.02 A),
since a search renumbers atoms when it maps the two ends.
"""
import csv, json, os, sys
from concurrent.futures import ProcessPoolExecutor
from pathlib import Path
import numpy as np

WS = None


def _init(ws):
    global WS
    WS = Path(ws)


def frames(p):
    L = open(p).read().splitlines(); out = []; i = 0
    while i < len(L):
        if not L[i].strip():
            i += 1; continue
        n = int(L[i]); out.append(np.array([[float(v) for v in l.split()[1:4]] for l in L[i + 2:i + 2 + n]])); i += n + 2
    return out


def spec(X):
    d = np.linalg.norm(X[:, None] - X[None], axis=-1)
    return np.sort(d[np.triu_indices(len(X), 1)])


def same(a, b):
    return len(a) == len(b) and float(np.abs(spec(a) - spec(b)).max()) < 0.02


def one(rid):
    run = WS / "search" / rid / "run"
    row = {"reaction_id": rid}
    try:
        S = frames(WS / "screen" / rid / "start.xyz")[0]; E = frames(WS / "screen" / rid / "end.xyz")[0]
        sp = frames(run / "conformers" / "start.xyz"); ep = frames(run / "conformers" / "end.xyz")
        pm = [p for p in json.loads((run / "pair_mechanisms.json").read_text()) if not p.get("not_in_this_run")]
        st = json.loads((run / "stats.json").read_text())
    except (OSError, ValueError, IndexError):
        return {**row, "readable": 0}
    ns = len(sp)
    si = [i for i, f in enumerate(sp) if same(f, S)]
    ei = [j for j, f in enumerate(ep) if same(f, E)]
    irc = [p["pair"] for p in pm if p["start_conformer"] in si and int(p["pair"].split("_")[2]) - ns in ei]
    tree = run / "pairs" / irc[0] / "tree" if irc else None
    done = bool(tree and (tree / "tree.json").exists())
    nodes = [f for f in os.listdir(tree) if f.startswith("node_") and f.endswith(".xyz") and "rejected" not in f] if done else []
    do = st.get("direct_only") or {}
    return {**row, "readable": 1, "n_pairs_possible": ns * len(ep), "n_pairs_queued": len(pm),
            "irc_ends_in_pools": int(bool(si) and bool(ei)), "irc_pair_queued": int(bool(irc)),
            "irc_pair_path_done": int(done), "irc_pair_path_split": int(len(nodes) > 1),
            "budget_reached": int(bool(st.get("search_budget_reached"))),
            "legs_not_run": int(do.get("legs_not_run", 0))}


if __name__ == "__main__":
    out = {r["reaction_id"]: r for r in csv.DictReader(open(sys.argv[2]))}
    ids = [i for i, r in out.items() if r["status"] in ("done", "no_route")]
    _init(sys.argv[1])
    with ProcessPoolExecutor(4, initializer=_init, initargs=(sys.argv[1],)) as pool:
        rows = list(pool.map(one, ids, chunksize=64))
    for r in rows:
        o = out[r["reaction_id"]]
        r["status"] = o["status"]; r["n_offtarget"] = ""
    res = {}
    for ln in open(WS / "search" / "results.jsonl"):
        if ln.strip():
            s = json.loads(ln); res.setdefault(s["id"], s) if s.get("status") == "error" else res.__setitem__(s["id"], s)
    for r in rows:
        r["n_offtarget"] = res.get(r["reaction_id"], {}).get("n_offtarget", "")
    fields = ["reaction_id", "status", "readable", "n_pairs_possible", "n_pairs_queued", "irc_ends_in_pools",
              "irc_pair_queued", "irc_pair_path_done", "irc_pair_path_split", "budget_reached", "legs_not_run", "n_offtarget"]
    with open(sys.argv[3] if len(sys.argv) > 3 else "search_diagnostics.csv", "w", newline="") as fh:
        w = csv.DictWriter(fh, fieldnames=fields, extrasaction="ignore"); w.writeheader(); w.writerows(rows)
    print(len(rows), "rows")