"""Build F0001's evidence tables from the raw screen output.

    python reproduce/build_tables.py --screen <screen results.jsonl> --reaction-qm <folder with the Reaction-QM files> [--out data]

Writes (into data/ by default):
  data/screen.csv.zip     one row per screened reaction (no geometries)
  data/population.csv.zip all 199,890 B3LYP-RXN reactions, flagged screened or not,
                          so the screened prefix can be compared with the whole set
  data/large/gxtb_ts.xyz.gz  the converged g-xTB saddles of matched reactions

Elements, total charge, multiplicity and fragment counts are read from the
reactant fragments in the h5 (the mapped SMILES are not reliable to parse for
atom counts).
"""

from __future__ import annotations

import argparse
import csv
import gzip
import io
import json
import zipfile
from contextlib import contextmanager
from pathlib import Path

HERE = Path(__file__).resolve().parent
DATA = HERE.parent / "data"

SYMBOLS = {1: "H", 6: "C", 7: "N", 8: "O", 9: "F", 14: "Si", 15: "P", 16: "S", 17: "Cl", 35: "Br", 53: "I"}


def gz_text(path: Path):
    """Byte-reproducible gzip (fixed mtime) so checksums only change with content."""
    return io.TextIOWrapper(gzip.GzipFile(path, "wb", compresslevel=9, mtime=0), encoding="utf-8", newline="")


@contextmanager
def zip_csv(path: Path):
    """Text handle whose contents land in a byte-reproducible zip (fixed timestamp) holding one
    file, path's name without .zip. Zip, not gzip: macOS Archive Utility misreads a gzipped CSV."""
    buf = io.StringIO(newline="")
    yield buf
    info = zipfile.ZipInfo(path.name.removesuffix(".zip"), date_time=(1980, 1, 1, 0, 0, 0))
    info.compress_type = zipfile.ZIP_DEFLATED
    info.external_attr = 0o644 << 16
    with zipfile.ZipFile(path, "w") as z:
        z.writestr(info, buf.getvalue().encode("utf-8"), compresslevel=9)


def element_key(elements: list[str]) -> str:
    """Hill-ordered element set, e.g. CHNO."""
    s = set(elements)
    head = [e for e in ("C", "H") if e in s]
    return "".join(head + sorted(s - set(head)))


def species_info(rqm_dir: Path, index: dict) -> dict:
    """Per reaction, from the reactant fragments in the h5: elements, total charge, multiplicity."""
    import h5py

    out = {}
    with h5py.File(rqm_dir / "B3LYPD3_TZVP.h5", "r") as h5:
        for rid, chunk in index.items():
            g = h5[chunk][rid]
            frags = sorted(k for k in g if k.startswith("R")) or ["TS"]
            z = [int(x) for k in frags for x in g[k]["atomic_numbers"][()]]
            out[rid] = {
                "elements": [SYMBOLS.get(n, str(n)) for n in z],
                "charge": sum(int(g[k]["charge"][()]) for k in frags),
                "multiplicity": int(g[frags[0]]["multiplicity"][()]),
                "n_reactant_frags": len(frags),
                # Some entries store a repeated reactant once (A + A kept as one A), so the
                # reactant side holds fewer atoms than the TS; flag rather than drop them.
                "balanced": int("TS" not in g or len(z) == len(g["TS"]["atomic_numbers"][()])),
                "ts_atoms": len(g["TS"]["atomic_numbers"][()]) if "TS" in g else len(z),
            }
    return out


def main() -> None:
    p = argparse.ArgumentParser()
    p.add_argument("--screen", type=Path, required=True)
    p.add_argument("--reaction-qm", type=Path, required=True)
    p.add_argument("--out", type=Path, default=DATA, help="where to write the tables (default data/)")
    a = p.parse_args()

    index = json.loads((a.reaction_qm / "B3LYPD3_TZVP_index.json").read_text())
    info = {}
    with open(a.reaction_qm / "B3LYPD3_TZVP_reaction_info.csv") as fh:
        for row in csv.DictReader(fh):
            info[row["reaction_id"]] = row
    spec = species_info(a.reaction_qm, index)

    recs = [json.loads(line) for line in a.screen.open()]
    screened = {r["reaction_id"]: r for r in recs}

    out = a.out
    out.mkdir(exist_ok=True)
    (out / "large").mkdir(exist_ok=True)

    pop_cols = ["reaction_id", "h5_chunk", "screened", "n_atoms", "n_heavy", "elements", "charge",
                "multiplicity", "n_reactant_frags", "balanced", "ref_dE_dagger_kcal", "ref_dE_kcal"]
    with zip_csv(out / "population.csv.zip") as fh:
        w = csv.writer(fh)
        w.writerow(pop_cols)
        for rid, chunk in index.items():
            row = info[rid]
            sp = spec[rid]
            els = sp["elements"]
            w.writerow([rid, chunk, int(rid in screened), sp["ts_atoms"], sum(e != "H" for e in els),
                        element_key(els), sp["charge"], sp["multiplicity"], sp["n_reactant_frags"], sp["balanced"],
                        f"{float(row['dE_dagger']):.4f}", f"{float(row['dE']):.4f}"])

    cols = ["reaction_id", "h5_chunk", "status", "detail", "direction", "n_atoms", "n_heavy", "elements",
            "charge", "multiplicity", "n_reactant_frags", "balanced", "ref_barrier_kcal", "gxtb_barrier_sep_kcal", "gxtb_barrier_complex_kcal",
            "ref_dE_kcal", "gxtb_dE_kcal", "t_ts", "t_irc", "t_total", "ref_reactant", "ref_product",
            "irc_a", "irc_b"]
    with zip_csv(out / "screen.csv.zip") as fh, \
            gz_text(out / "large" / "gxtb_ts.xyz.gz") as xyz:
        w = csv.writer(fh)
        w.writerow(cols)
        for r in sorted(recs, key=lambda r: r["reaction_id"]):
            rid = r["reaction_id"]
            sp = spec[rid]
            row = dict(r, h5_chunk=index[rid], elements=element_key(sp["elements"]), charge=sp["charge"],
                       multiplicity=sp["multiplicity"], n_reactant_frags=sp["n_reactant_frags"], balanced=sp["balanced"], ref_dE_kcal=f"{float(info[rid]['dE']):.4f}")
            w.writerow(["" if row.get(c) is None else row.get(c, "") for c in cols])
            if r.get("ts_xyz"):
                lines = r["ts_xyz"].splitlines()
                xyz.write(f"{len(lines)}\n{rid} g-xTB saddle (IRC-matched, {r.get('direction')})\n")
                xyz.write("\n".join(lines) + "\n")
    print("wrote", *(sorted(str(q) for q in out.rglob("*.gz"))))


if __name__ == "__main__":
    main()