"""Batch driver: run `mepd channels` over a table of benchmark reactions.

    python -m benchmark.reaction_qm_channels --dataset SET.csv --inputs mepd_profile.toml --output RUNS --dry-run --limit 3
    python -m benchmark.reaction_qm_channels --dataset SET.csv --inputs mepd_profile.toml --output RUNS --jobs 4
    python -m benchmark.reaction_qm_channels --output RUNS --summarize

Run from the folder above src/, with PYTHONPATH=src and mepd on PATH. Each row
of the table names start, end and reference-TS xyz files under --geom-root
(default data/reaction_qm/). Arguments after `--` go to `mepd channels`.

* Endpoints are xyz geometries, not SMILES, so nothing is re-embedded.
* The two endpoint files need not share an atom ordering; mepd's atom mapping
  reconciles them (turn it off with --no-atom-mapping when they do share one).
* Every run directory gets the reference TS and its barrier in
  `_driver_status.json`, so a found TS can be scored (see score_channels).

Runs are resumable (finished reactions are skipped), each reaction has a
wall-clock timeout, and --summarize tabulates the results.
"""

from __future__ import annotations

import argparse
import csv
import dataclasses
import json
import os
import shlex
import signal
import subprocess
import sys
import time
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
from typing import Iterable, Optional

__all__ = ["Reaction", "load_reactions", "build_command", "run_reaction", "run_all", "summarize"]

_HERE = Path(__file__).resolve().parent
_REPO_ROOT = _HERE.parent.parent

DEFAULT_GEOM_ROOT = _REPO_ROOT / "data" / "reaction_qm"

DEFAULT_N_CONFORMERS = 0  # mepd defaults: no cap, selection by --pairs-per-mechanism
DEFAULT_MAX_PAIRS = 0
DEFAULT_TIMEOUT_S = 2 * 3600


@dataclasses.dataclass(frozen=True)
class Reaction:
    idx: int
    reaction_id: str
    n_heavy: int
    n_atoms: int
    charge: int
    multiplicity: int
    n_reactant_frags: int
    n_product_frags: int
    elements: str
    barrier_kcal: float
    start_smiles: str
    end_smiles: str
    dataset_start_smiles: str
    dataset_end_smiles: str
    perception_agrees: int
    start_xyz: str
    end_xyz: str
    ts_reference_xyz: str
    # Present when the reference TS is a g-xTB saddle and `barrier_kcal` is
    # its g-xTB barrier.
    ref_b3lyp_barrier_kcal: Optional[float] = None
    barrier_delta_kcal: Optional[float] = None

    @property
    def slug(self) -> str:
        return f"rxn_{self.idx:03d}_{self.reaction_id}"

    def start_path(self, geom_root: Path) -> Path:
        return Path(geom_root) / self.start_xyz

    def end_path(self, geom_root: Path) -> Path:
        return Path(geom_root) / self.end_xyz

    def ts_path(self, geom_root: Path) -> Path:
        return Path(geom_root) / self.ts_reference_xyz


def load_reactions(dataset: Path) -> list:
    dataset = Path(dataset)
    if not dataset.is_file():
        raise FileNotFoundError(
            f"benchmark CSV not found: {dataset}"
        )
    int_fields = {
        "idx", "n_heavy", "n_atoms", "charge", "multiplicity",
        "n_reactant_frags", "n_product_frags", "perception_agrees",
    }
    float_fields = {"barrier_kcal", "ref_b3lyp_barrier_kcal", "barrier_delta_kcal"}
    known = {f.name for f in dataclasses.fields(Reaction)}
    out = []
    with dataset.open(newline="") as fh:
        for row in csv.DictReader(fh):
            kw = {}
            for k, v in row.items():
                # Tolerate extra columns rather than crashing: different
                # benchmark tables carry different extra columns.
                if k not in known:
                    continue
                if k in int_fields:
                    kw[k] = int(v)
                elif k in float_fields:
                    kw[k] = float(v) if v not in ("", None) else None
                else:
                    kw[k] = v
            out.append(Reaction(**kw))
    return out


def select(
    reactions: list,
    *,
    limit: Optional[int] = None,
    sort_by_size: bool = False,
    max_atoms: Optional[int] = None,
    min_barrier: Optional[float] = None,
    max_barrier: Optional[float] = None,
    molecularity: Optional[int] = None,
    only: Optional[Iterable[str]] = None,
) -> list:
    out = list(reactions)
    if only:
        wanted = set(only)
        out = [r for r in out if r.reaction_id in wanted or str(r.idx) in wanted]
    if max_atoms is not None:
        out = [r for r in out if r.n_atoms <= max_atoms]
    if min_barrier is not None:
        out = [r for r in out if r.barrier_kcal >= min_barrier]
    if max_barrier is not None:
        out = [r for r in out if r.barrier_kcal <= max_barrier]
    if molecularity is not None:
        out = [r for r in out if r.n_reactant_frags == molecularity]
    if sort_by_size:
        out.sort(key=lambda r: (r.n_atoms, r.n_heavy))
    if limit is not None:
        out = out[:limit]
    return out


def build_command(
    rxn: Reaction,
    outdir: Path,
    *,
    geom_root: Path = DEFAULT_GEOM_ROOT,
    inputs: Path,
    mepd: str = "mepd",
    atom_mapping: bool = True,
    backend: str = "rdkit",
    crest_method: Optional[str] = None,
    crest_threads: Optional[int] = None,
    crest_ewin: Optional[float] = None,
    rdkit_ewin: Optional[float] = None,
    n_embed: Optional[int] = None,
    conformers_only: bool = False,
    n_conformers: int = DEFAULT_N_CONFORMERS,
    max_pairs: int = DEFAULT_MAX_PAIRS,
    workers: Optional[int] = None,
    random_seed: int = 0,
    hessian: bool = True,
    parallel: bool = False,
    parallel_workers: Optional[int] = None,
    debug_dump: bool = False,
    atom_mapping_metric: Optional[str] = None,
    atom_mapping_candidates: Optional[int] = None,
    extra_args: Iterable[str] = (),
) -> list:
    start, end = rxn.start_path(geom_root), rxn.end_path(geom_root)
    for tag, fp in (("--start", start), ("--end", end)):
        if not fp.is_file():
            raise FileNotFoundError(
                f"{rxn.slug}: {tag} geometry missing: {fp}"
            )
    cmd = [
        mepd, "channels",
        "--start", str(start),
        "--end", str(end),
        "--inputs", str(inputs),
        "--method", "conformers",
        "--backend", backend,
        "--n-conformers", str(n_conformers),
        "--max-pairs", str(max_pairs),
        "--random-seed", str(random_seed),
        "--charge", str(rxn.charge),
        "--multiplicity", str(rxn.multiplicity),
        "--output", str(outdir),
    ]
    if backend == "crest":
        if crest_method:
            cmd += ["--crest-method", crest_method]
        if crest_threads is not None:
            cmd += ["--crest-threads", str(crest_threads)]
        if crest_ewin is not None:
            cmd += ["--crest-ewin", str(crest_ewin)]
    if backend == "rdkit":
        if rdkit_ewin is not None:
            cmd += ["--rdkit-ewin", str(rdkit_ewin)]
        if n_embed is not None:
            cmd += ["--n-embed", str(n_embed)]
    if conformers_only:
        cmd.append("--conformers-only")
    cmd.append("--atom-mapping" if atom_mapping else "--no-atom-mapping")
    if workers is not None:
        cmd += ["--workers", str(workers)]
    if atom_mapping_metric:
        cmd += ["--atom-mapping-metric", atom_mapping_metric]
    if atom_mapping_candidates is not None:
        cmd += ["--atom-mapping-candidates", str(atom_mapping_candidates)]
    cmd.append("--validate-minima-with-hessian" if hessian else "--no-validate-minima-with-hessian")
    if parallel:
        cmd.append("--parallel")
        if parallel_workers is not None:
            cmd += ["--parallel-workers", str(parallel_workers)]
    if debug_dump:
        cmd.append("--debug-dump")
    cmd += list(extra_args)
    return cmd


def _job_env(threads: int) -> dict:
    """Environment for one reaction's subprocess.

    `gxtb-xtb` is OpenMP-threaded and, with no OMP_NUM_THREADS set, uses every
    core. For molecules this small the threads do not pay for themselves, so
    parallelism belongs across reactions: each job gets `threads` threads, so
    that many parallel jobs do not oversubscribe the machine."""
    env = dict(os.environ)
    for var in ("OMP_NUM_THREADS", "MKL_NUM_THREADS", "OPENBLAS_NUM_THREADS",
                "NUMEXPR_NUM_THREADS", "VECLIB_MAXIMUM_THREADS"):
        env[var] = str(threads)
    return env


def _status_path(outdir: Path) -> Path:
    return outdir / "_driver_status.json"


def is_done(outdir: Path) -> bool:
    st = _status_path(outdir)
    if not st.is_file():
        return False
    try:
        return json.loads(st.read_text()).get("returncode") == 0
    except Exception:
        return False


def _inspect_output(outdir: Path) -> dict:
    """Count what a finished channels run produced, in the vocabulary the
    fixed classifier uses: a direct channel, a multistep alternate channel,
    or an off-target exit."""
    def _n(sub: str) -> int:
        d = outdir / sub
        return len([p for p in d.iterdir() if p.is_dir()]) if d.is_dir() else 0

    out = {
        "n_channels": _n("channels"),
        "n_alternate_channels": _n("alternate-channels"),
        "n_offtarget_exit_channels": _n("offtarget-exit-channels"),
        "n_pairs": _n("pairs"),
        "has_network": (outdir / "network.json").is_file(),
    }
    stats_fp = outdir / "stats.json"
    if stats_fp.is_file():
        try:
            out["stats"] = json.loads(stats_fp.read_text())
        except Exception:
            pass
    log = outdir / "channels.log"
    if log.is_file():
        try:
            tail = log.read_text(errors="replace").splitlines()
            out["log_tail"] = "\n".join(tail[-12:])
            for line in tail:
                if line.startswith("--atom-mapping: reindex"):
                    out["atom_mapping_reindexed"] = True
                elif line.startswith("--atom-mapping: keeping"):
                    out["atom_mapping_reindexed"] = False
        except Exception:
            pass
    return out


def run_reaction(
    rxn: Reaction,
    root: Path,
    *,
    geom_root: Path = DEFAULT_GEOM_ROOT,
    timeout: int = DEFAULT_TIMEOUT_S,
    resume: bool = True,
    dry_run: bool = False,
    threads_per_job: int = 1,
    **cmd_kwargs,
) -> dict:
    outdir = Path(root) / rxn.slug
    if resume and is_done(outdir):
        rec = json.loads(_status_path(outdir).read_text())
        rec["skipped"] = "already complete"
        return rec

    cmd = build_command(rxn, outdir, geom_root=geom_root, **cmd_kwargs)
    if dry_run:
        return {"idx": rxn.idx, "reaction_id": rxn.reaction_id, "dry_run": True,
                "command": " ".join(shlex.quote(c) for c in cmd)}

    outdir.mkdir(parents=True, exist_ok=True)
    log = outdir / "channels.log"
    started = time.time()
    rec: dict = {
        "idx": rxn.idx,
        "reaction_id": rxn.reaction_id,
        "n_heavy": rxn.n_heavy,
        "n_atoms": rxn.n_atoms,
        "n_reactant_frags": rxn.n_reactant_frags,
        "elements": rxn.elements,
        "reference_barrier_kcal": rxn.barrier_kcal,
        "reference_ts_xyz": str(rxn.ts_path(geom_root)),
        "start_smiles": rxn.start_smiles,
        "end_smiles": rxn.end_smiles,
        "command": " ".join(shlex.quote(c) for c in cmd),
        "started_at": time.strftime("%Y-%m-%dT%H:%M:%S", time.localtime(started)),
    }
    try:
        with log.open("w") as fh:
            fh.write(rec["command"] + "\n\n")
            fh.flush()
            proc = subprocess.Popen(
                cmd, stdout=fh, stderr=subprocess.STDOUT,
                cwd=str(root), start_new_session=True,
                env=_job_env(threads_per_job),
            )
            try:
                rec["returncode"] = proc.wait(timeout=timeout or None)
            except subprocess.TimeoutExpired:
                try:
                    os.killpg(os.getpgid(proc.pid), signal.SIGTERM)
                    proc.wait(timeout=60)
                except Exception:
                    try:
                        os.killpg(os.getpgid(proc.pid), signal.SIGKILL)
                    except Exception:
                        pass
                rec["returncode"] = -1
                rec["error"] = f"timeout after {timeout}s"
    except FileNotFoundError as exc:
        rec["returncode"] = -1
        rec["error"] = f"{exc} (is `mepd` on PATH?)"
    rec["wall_seconds"] = round(time.time() - started, 1)
    rec.update(_inspect_output(outdir))
    _status_path(outdir).write_text(json.dumps(rec, indent=2))
    return rec


def _report(i: int, n: int, rec: dict) -> None:
    if rec.get("dry_run"):
        print(f"[{i}/{n}] {rec['command']}")
        return
    if rec.get("skipped"):
        print(f"[{i}/{n}] rxn_{rec['idx']:03d} {rec['reaction_id']}: skipped ({rec['skipped']})")
        return
    ok = "ok " if rec.get("returncode") == 0 else "FAIL"
    print(
        f"[{i}/{n}] rxn_{rec['idx']:03d} {rec['reaction_id']}: {ok} "
        f"rc={rec.get('returncode')} {rec.get('wall_seconds')}s "
        f"channels={rec.get('n_channels')} "
        f"altchan={rec.get('n_alternate_channels')} "
        f"offtarget={rec.get('n_offtarget_exit_channels')}"
        + (f" -- {rec['error']}" if rec.get("error") else "")
    )
    sys.stdout.flush()


def run_all(
    reactions: list,
    root: Path,
    *,
    jobs: int = 1,
    timeout: int = DEFAULT_TIMEOUT_S,
    resume: bool = True,
    dry_run: bool = False,
    threads_per_job: int = 1,
    **cmd_kwargs,
) -> list:
    root = Path(root)
    root.mkdir(parents=True, exist_ok=True)
    n = len(reactions)
    results: list = []

    def _one(rxn):
        return run_reaction(rxn, root, timeout=timeout, resume=resume,
                            dry_run=dry_run, threads_per_job=threads_per_job,
                            **cmd_kwargs)

    if jobs <= 1 or dry_run:
        for i, rxn in enumerate(reactions, 1):
            rec = _one(rxn)
            results.append(rec)
            _report(i, n, rec)
    else:
        with ThreadPoolExecutor(max_workers=jobs) as pool:
            for i, rec in enumerate(pool.map(_one, reactions), 1):
                results.append(rec)
                _report(i, n, rec)

    if not dry_run:
        ok = sum(1 for r in results if r.get("returncode") == 0 or r.get("skipped"))
        print(f"\ndone: {ok}/{n} succeeded")
    return results


def summarize(root: Path) -> list:
    root = Path(root)
    recs = []
    for d in sorted(root.glob("rxn_*")):
        st = _status_path(d)
        if st.is_file():
            try:
                recs.append(json.loads(st.read_text()))
            except Exception:
                pass
    if not recs:
        print(f"no completed runs under {root}")
        return recs
    ok = [r for r in recs if r.get("returncode") == 0]
    connected = [r for r in ok if r.get("n_channels") or r.get("n_alternate_channels")]
    print(f"{len(recs)} run(s) under {root}: {len(ok)} exited 0, "
          f"{len(connected)} connected --start to --end (direct or multistep)")
    print(f"{'idx':>4} {'reaction_id':<18} {'atoms':>5} {'frags':>5} {'ref_kcal':>9} "
          f"{'rc':>3} {'wall_s':>9} {'chan':>5} {'altch':>5} {'offtg':>5}  aam")
    for r in sorted(recs, key=lambda r: r["idx"]):
        aam = {True: "reindexed", False: "kept", None: "-"}[r.get("atom_mapping_reindexed")]
        print(f"{r['idx']:>4} {r.get('reaction_id',''):<18} {r.get('n_atoms','?'):>5} "
              f"{r.get('n_reactant_frags','?'):>5} {r.get('reference_barrier_kcal','?'):>9} "
              f"{r.get('returncode','?'):>3} {r.get('wall_seconds','?'):>9} "
              f"{r.get('n_channels',0):>5} {r.get('n_alternate_channels',0):>5} "
              f"{r.get('n_offtarget_exit_channels',0):>5}  {aam}")
    tot = sum(r.get("wall_seconds") or 0 for r in recs)
    print(f"total wall time across runs: {tot/3600:.2f} h")
    (root / "sweep_summary.json").write_text(json.dumps(recs, indent=2))
    return recs


def _parse_args(argv=None):
    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("--dataset", type=Path)
    p.add_argument("--geom-root", type=Path, default=DEFAULT_GEOM_ROOT)
    p.add_argument("--inputs", type=Path)
    p.add_argument("--output", type=Path, required=True)
    p.add_argument("--mepd", default="mepd")
    p.add_argument("--limit", type=int)
    p.add_argument("--only", nargs="+", help="reaction ids or idx values")
    p.add_argument("--sort-by-size", action="store_true")
    p.add_argument("--max-atoms", type=int)
    p.add_argument("--min-barrier", type=float)
    p.add_argument("--max-barrier", type=float)
    p.add_argument("--molecularity", type=int, help="1=unimolecular, 2=bimolecular, ...")
    p.add_argument("--backend", choices=("rdkit", "crest"), default="rdkit",
                   help="`mepd channels --backend`: conformer generator")
    p.add_argument("--crest-method", help="`mepd channels --crest-method`, e.g. --gfn2")
    p.add_argument("--crest-threads", type=int,
                   help="`mepd channels --crest-threads` (CREST's -T)")
    p.add_argument("--crest-ewin", type=float, help="`mepd channels --crest-ewin` (kcal/mol)")
    p.add_argument("--rdkit-ewin", type=float, help="`mepd channels --rdkit-ewin` (kcal/mol, MMFF)")
    p.add_argument("--n-embed", type=int, help="`mepd channels --n-embed` (0 = auto)")
    p.add_argument("--conformers-only", action="store_true",
                   help="`mepd channels --conformers-only`: size pools/pairs, no path search")
    p.add_argument("--n-conformers", type=int, default=DEFAULT_N_CONFORMERS,
                   help="0 = no cap")
    p.add_argument("--max-pairs", type=int, default=DEFAULT_MAX_PAIRS, help="0 = no cap")
    p.add_argument("--workers", type=int,
                   help="`mepd channels --workers`: pair searches / TS opts in parallel per reaction")
    p.add_argument("--random-seed", type=int, default=0)
    p.add_argument("--no-atom-mapping", action="store_true")
    p.add_argument("--atom-mapping-metric")
    p.add_argument("--atom-mapping-candidates", type=int)
    p.add_argument("--no-hessian", action="store_true")
    p.add_argument("--parallel", action="store_true")
    p.add_argument("--parallel-workers", type=int)
    p.add_argument("--debug-dump", action="store_true")
    p.add_argument("--jobs", type=int, default=1,
                   help="reactions to run concurrently (see --threads-per-job)")
    p.add_argument("--threads-per-job", type=int, default=1,
                   help="OMP/BLAS threads per reaction; 1 is right for this set")
    p.add_argument("--timeout", type=int, default=DEFAULT_TIMEOUT_S, help="seconds per reaction; 0 = no limit")
    p.add_argument("--no-resume", action="store_true")
    p.add_argument("--dry-run", action="store_true")
    p.add_argument("--summarize", action="store_true")
    p.add_argument("extra", nargs="*", help="extra args passed through to `mepd channels`")
    return p.parse_args(argv)


def main(argv=None) -> int:
    args = _parse_args(argv)
    # Resolve every path before use. `run_reaction` launches mepd with
    # cwd=<output root>, so a *relative* --output is re-resolved against that
    # root by the child and the whole tree lands one level too deep
    # (RUNS/RUNS/rxn_*), leaving the driver inspecting
    # empty directories and reporting zero channels everywhere.
    args.output = args.output.resolve()
    args.geom_root = args.geom_root.resolve()
    if not args.summarize and (args.dataset is None or args.inputs is None):
        raise SystemExit("--dataset and --inputs are required (except with --summarize)")
    args.inputs = args.inputs.resolve() if args.inputs else None
    args.dataset = args.dataset.resolve() if args.dataset else None
    if args.summarize:
        summarize(args.output)
        return 0

    reactions = select(
        load_reactions(args.dataset),
        limit=args.limit, sort_by_size=args.sort_by_size, max_atoms=args.max_atoms,
        min_barrier=args.min_barrier, max_barrier=args.max_barrier,
        molecularity=args.molecularity, only=args.only,
    )
    if not reactions:
        print("no reactions selected")
        return 1

    n_conf = args.n_conformers
    max_pairs = args.max_pairs
    print(f"{len(reactions)} reaction(s); backend={args.backend} "
          f"n_conformers={n_conf} max_pairs={max_pairs} "
          f"hessian={not args.no_hessian} atom_mapping={not args.no_atom_mapping} "
          f"jobs={args.jobs}x{args.threads_per_job} thread(s)")
    print(f"inputs={args.inputs}  output={args.output}  geometries={args.geom_root}\n")

    run_all(
        reactions, args.output,
        jobs=args.jobs, timeout=args.timeout, resume=not args.no_resume,
        threads_per_job=args.threads_per_job,
        dry_run=args.dry_run, geom_root=args.geom_root, inputs=args.inputs,
        mepd=args.mepd, atom_mapping=not args.no_atom_mapping,
        backend=args.backend, crest_method=args.crest_method,
        crest_threads=args.crest_threads, crest_ewin=args.crest_ewin,
        rdkit_ewin=args.rdkit_ewin, n_embed=args.n_embed,
        conformers_only=args.conformers_only,
        n_conformers=n_conf, max_pairs=max_pairs, workers=args.workers,
        random_seed=args.random_seed,
        hessian=not args.no_hessian, parallel=args.parallel,
        parallel_workers=args.parallel_workers, debug_dump=args.debug_dump,
        atom_mapping_metric=args.atom_mapping_metric,
        atom_mapping_candidates=args.atom_mapping_candidates,
        extra_args=args.extra,
    )
    return 0


if __name__ == "__main__":
    raise SystemExit(main())