"""Build data/ from an explorer run and the DFT single points.

    python build_tables.py --run-dir RUN --dft-dir DFT --out DATA

RUN: the `mepd discovery expand` output folder (summary.json, species/, ts/, channels/).
DFT: ts_sp.json and r_sp.json from dft_single_points.py (wB97X-D/def2-TZVP energies, Hartree).
Needs mepd (g-xTB energies of the separated reactants and the TSs) and RDKit.
"""
import argparse
import csv
import json
import os
from pathlib import Path

import numpy as np

HARTREE_KCAL = 627.5095
HARTREE_KCAL_RUN = 627.5094740631   # the explorer's own conversion, for energies relative to its species
CLOSED_KCAL = 300.0                 # the kinetics close a step out of reach by raising its TS by 1 Hartree
BOHR_PER_ANGSTROM = 1.8897259886


def read_xyz(fp):
    lines = Path(fp).read_text().split("\n")
    n = int(lines[0])
    rows = [line.split() for line in lines[2:2 + n]]
    return [r[0] for r in rows], np.array([[float(v) for v in r[1:4]] for r in rows])


def stereo(fp):
    """endo or exo for a 5-norbornene-2-carboxylate: the ester carbon nearer the C=C than the CH2 bridge = endo."""
    from rdkit import Chem
    from rdkit.Chem import rdDetermineBonds

    m = Chem.MolFromXYZFile(str(fp))
    rdDetermineBonds.DetermineBonds(m, charge=0)
    x = m.GetConformer().GetPositions()
    cc = [a for b in m.GetBonds() if b.GetBondType() == Chem.BondType.DOUBLE and b.IsInRing()
          for a in (b.GetBeginAtomIdx(), b.GetEndAtomIdx())]
    carbons = lambda a: {n.GetIdx() for n in m.GetAtomWithIdx(a).GetNeighbors() if n.GetSymbol() == "C"}
    bridge = [a.GetIdx() for a in m.GetAtoms() if a.IsInRing() and a.GetSymbol() == "C"
              and a.GetTotalNumHs(includeNeighbors=True) == 2 and all(carbons(n) & set(cc) for n in carbons(a.GetIdx()))]
    ester = [a.GetIdx() for a in m.GetAtoms() if a.GetSymbol() == "C" and not a.IsInRing() and a.GetDegree() == 3]
    if len(cc) != 2 or len(bridge) != 1 or len(ester) != 1:
        return ""
    e = x[ester[0]]
    return "endo" if np.linalg.norm(e - x[cc].mean(0)) < np.linalg.norm(e - x[bridge[0]]) else "exo"


def step_kind(label):
    if label.startswith("fast_"):
        return "proposal"
    for key, kind in (("offtarget", "exit"), ("alternate", "alternate"), ("_channels_channel_", "channel")):
        if key in label:
            return kind
    return ""


def write_steps(s, run_dir, out):
    """steps.csv (every verified step) and network.xyz (every species and transition state)."""
    sp = s["species"]
    e0, rel0 = sp[0]["energy"], sp[0]["rel_energy_kcal"]
    open_ = [st["kinetics_barrier_kcal"][0] - st["barrier_kcal"][0] < CLOSED_KCAL for st in s["steps"]]
    # free energies relative to the separated reactants, through the steps the kinetics kept open
    g = {0: 0.0}
    grew = True
    while grew:
        grew = False
        for st, o in zip(s["steps"], open_):
            fwd, rev = st["kinetics_barrier_kcal"]
            if o and (st["a"] in g) != (st["b"] in g):
                if st["a"] in g:
                    g[st["b"]] = g[st["a"]] + fwd - rev
                else:
                    g[st["a"]] = g[st["b"]] - fwd + rev
                grew = True
    run_file = lambda fp: run_dir / Path(fp).relative_to(Path(fp).parts[0])
    with open(out / "steps.csv", "w", newline="") as fh:
        w = csv.writer(fh, lineterminator="\n")
        w.writerow(["step", "from_species", "to_species", "kind", "label", "ts_rel_energy_kcal", "barrier_fwd_kcal",
                    "barrier_rev_kcal", "ts_dG_kcal", "dG_barrier_fwd_kcal", "dG_barrier_rev_kcal", "in_reach"])
        for i, (st, o) in enumerate(zip(s["steps"], open_)):
            fwd, rev = st["barrier_kcal"]
            gf, gr = st["kinetics_barrier_kcal"]
            free = [f"{g[st['a']] + gf:.3f}", f"{gf:.3f}", f"{gr:.3f}"] if o else ["", "", ""]
            w.writerow([i, st["a"], st["b"], step_kind(st["label"]), st["label"],
                        f"{rel0 + (st['ts_energy'] - e0) * HARTREE_KCAL_RUN:.3f}", f"{fwd:.3f}", f"{rev:.3f}",
                        *free, int(o)])
    frames = [(f"species {k}", run_dir / "species" / os.path.basename(x["file"])) for k, x in enumerate(sp)]
    frames += [(f"step {i} transition state", run_file(st["files"]["ts"])) for i, st in enumerate(s["steps"])]
    with open(out / "network.xyz", "w") as fh:
        for name, fp in frames:
            sy, x = read_xyz(fp)
            fh.write(f"{len(sy)}\n{name}\n" + "".join(f"{a:2s} {p[0]:15.8f} {p[1]:15.8f} {p[2]:15.8f}\n"
                                                         for a, p in zip(sy, x)))


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--run-dir", required=True, type=Path)
    ap.add_argument("--dft-dir", required=True, type=Path)
    ap.add_argument("--out", required=True, type=Path)
    a = ap.parse_args()
    s = json.loads((a.run_dir / "summary.json").read_text())
    kin = s["rounds"][-1]["kinetics"]
    conc = kin["final_concentration"]
    species_file = lambda k: a.run_dir / "species" / os.path.basename(s["species"][k]["file"])
    a.out.mkdir(parents=True, exist_ok=True)
    write_steps(s, a.run_dir, a.out)

    adducts = {}
    with open(a.out / "species.csv", "w", newline="") as fh:
        w = csv.writer(fh, lineterminator="\n")
        w.writerow(["species", "smiles", "stereo", "rel_energy_kcal", "final_fraction"])
        for k, sp in enumerate(s["species"]):
            st = stereo(species_file(k)) if sp["smiles"] == "COC(=O)C1CC2C=CC1C2" else ""
            if st:
                adducts[k] = st
            w.writerow([k, sp["smiles"], st, f"{sp['rel_energy_kcal']:.3f}", f"{conc[k]:.6g}"])

    import mepd
    from qcdata import Structure

    from mepd.discovery.free_energy import molecules
    from mepd.engines.gxtb import GXTBCalculator
    from mepd.nodes.node import StructureNode

    eng = GXTBCalculator()
    node = lambda sy, x: StructureNode(structure=Structure(symbols=list(sy), geometry=np.asarray(x) * BOHR_PER_ANGSTROM))
    sy, x = read_xyz(species_file(0))
    e_sep = sum(float(mepd.energies([node([sy[i] for i in p], x[p])], eng)[0]) for p in molecules(sy, x))
    dft_ts = json.loads((a.dft_dir / "ts_sp.json").read_text())
    dft_r = list(json.loads((a.dft_dir / "r_sp.json").read_text()).values())[0]

    with open(a.out / "cycloaddition_ts.csv", "w", newline="") as fh:
        w = csv.writer(fh, lineterminator="\n")
        w.writerow(["adduct_species", "stereo", "ts_file", "gxtb_ts_hartree", "gxtb_reactants_separated_hartree",
                    "gxtb_dG_kcal", "dft_ts_hartree", "dft_reactants_separated_hartree"])
        for st in s["steps"]:
            if not st["label"].startswith("ch_") or 0 not in (st["a"], st["b"]):
                continue
            k = st["b"] if st["a"] == 0 else st["a"]
            if k not in adducts:
                continue
            dg = st["kinetics_barrier_kcal"][0 if st["a"] == 0 else 1]
            ts_fp = a.run_dir / Path(st["files"]["ts"]).relative_to(Path(st["files"]["ts"]).parts[0])
            e_ts = float(mepd.energies([node(*read_xyz(ts_fp))], eng)[0])
            w.writerow([k, adducts[k], f"ts_{k}.xyz", f"{e_ts:.10f}", f"{e_sep:.10f}", f"{dg:.4f}",
                        f"{dft_ts[f'ts_{k}.xyz']:.10f}", f"{dft_r:.10f}"])


if __name__ == "__main__":
    main()