"""Build F0003's tables from the scored channels sweeps.

    python reproduce/build_tables.py [--study reproduce/rerun] [--out data]

Inputs (in the study folder):
  reaction_qm_scores.json             score_channels output, Reaction-QM 500 set
  geodesic_combined_scores.json       best of the mapped and unmapped Geodesic searches
  sets/{reaction_qm_500,geodesic_66}.csv   the benchmark sets
  reaction_qm_runs/, geodesic_runs_mapped/, geodesic_runs_unmapped/   run folders (TS geometries)
  data/reaction_qm/                   reference g-xTB TS geometries

Writes data/scores.csv (one row per reaction) and data/lower_ts.xyz.gz
(for every reaction with a lower TS: the reference g-xTB TS, then each TS of the
best route found, comment line = "<reaction_id> <role> E=<hartree>").
"""

from __future__ import annotations

import argparse
import csv
import gzip
import io
import json
from pathlib import Path

HARTREE_TO_KCAL = 627.509474
LOWER_KCAL = 1.0  # strictly more than this below the reference counts as lower
HERE = Path(__file__).resolve().parent
DATA = HERE.parent / "data"

SETS = {
    "reaction_qm": ("reaction_qm_scores.json", "sets/reaction_qm_500.csv", {None: "reaction_qm_runs"}),
    "geodesic": ("geodesic_combined_scores.json", "sets/geodesic_66.csv",
                 {"geodesic_runs_mapped": "geodesic_runs_mapped", "geodesic_runs_unmapped": "geodesic_runs_unmapped"}),
}


def outcome(r: dict) -> str:
    if not r.get("n_found"):
        return "no_channel"
    if r["improvement_kcal"] > LOWER_KCAL:
        return "lower"
    if r["improvement_kcal"] < -LOWER_KCAL:
        return "higher"
    return "matched"


def ts_files(run_dir: Path, kind: str, group: str) -> list[Path]:
    if kind == "channel":
        return [run_dir / "channels" / group / "ts.xyz"]
    return sorted((run_dir / "alternate-channels" / group).glob("step_*/ts.xyz"))


def read_xyz(p: Path) -> list[str]:
    lines = p.read_text().splitlines()
    n = int(lines[0])
    return lines[2: 2 + n]


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--study", type=Path, default=HERE / "rerun",
                    help="folder with the scored searches, laid out as rerun/ (default)")
    ap.add_argument("--out", type=Path, default=DATA, help="where to write the two tables (default data/)")
    a = ap.parse_args()
    pg = a.study.expanduser()

    rows, frames = [], []
    for set_name, (scores_fn, set_csv, roots) in SETS.items():
        scores = json.loads((pg / scores_fn).read_text())
        bench = {r["reaction_id"]: r for r in csv.DictReader(open(pg / set_csv))}
        for r in scores:
            b = bench[r["reaction_id"]]
            o = outcome(r)
            floor = r["reactant_floor_hartree"]
            row = {
                "set": set_name, "reaction_id": r["reaction_id"], "n_atoms": b["n_atoms"], "n_heavy": b["n_heavy"],
                "elements": b["elements"], "charge": b["charge"], "multiplicity": b["multiplicity"], "n_reactant_frags": b["n_reactant_frags"], "batch": b.get("batch", ""),
                "returncode": r["returncode"], "wall_seconds": r["wall_seconds"], "outcome": o,
                "n_routes_found": r["n_found"], "best_kind": r.get("best_kind", ""),
                "ref_ts_hartree": f"{floor + r['ref_barrier_kcal'] / HARTREE_TO_KCAL:.8f}",
                "best_ts_hartree": (f"{floor + r['best_barrier_kcal'] / HARTREE_TO_KCAL:.8f}" if r.get("n_found") else ""),
                "lower_by_kcal": r.get("improvement_kcal", ""),
                "excluded_routes": len(r.get("excluded") or []),
            }
            rows.append(row)
            if o != "lower":
                continue
            ref = pg / "data" / "reaction_qm" / b["ts_reference_xyz"]
            frames.append((f"{r['reaction_id']} reference E={row['ref_ts_hartree']}", read_xyz(ref)))
            best = min((f for f in r["found"]), key=lambda f: f["barrier_kcal"])
            run_dir = next((pg / roots[best.get("root")]).glob(f"rxn_*_{r['reaction_id']}"))
            files = ts_files(run_dir, best["kind"], best["group"])
            for i, fp in enumerate(files):
                role = "found" if best["kind"] == "channel" else f"found_step{i + 1}of{len(files)}"
                frames.append((f"{r['reaction_id']} {role}", read_xyz(fp)))

    a.out.mkdir(exist_ok=True)
    with open(a.out / "scores.csv", "w", encoding="utf-8", newline="") as out:
        w = csv.DictWriter(out, fieldnames=list(rows[0]))
        w.writeheader()
        w.writerows(rows)
    with io.TextIOWrapper(gzip.GzipFile(a.out / "lower_ts.xyz.gz", "wb", compresslevel=9, mtime=0),
                          encoding="utf-8") as out:
        for comment, atoms in frames:
            out.write(f"{len(atoms)}\n{comment}\n" + "\n".join(atoms) + "\n")
    print(f"{len(rows)} rows; {len(frames)} TS frames")


if __name__ == "__main__":
    main()