"""Write the structures for the DFT single points from an exploration: the starting reactants and the
transition state of each cycloaddition step that channel sampling found.

    python collect_dft_inputs.py --run-dir RUN --out DIR
"""
import argparse
import json
import os
import shutil
from pathlib import Path

ADDUCT = "COC(=O)C1CC2C=CC1C2"

ap = argparse.ArgumentParser()
ap.add_argument("--run-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())
a.out.mkdir(parents=True, exist_ok=True)
shutil.copy(a.run_dir / "species" / os.path.basename(s["species"][0]["file"]), a.out / "reactants.xyz")
for st in s["steps"]:
    if st["label"].startswith("ch_") and 0 in (st["a"], st["b"]):
        k = st["b"] if st["a"] == 0 else st["a"]
        if s["species"][k]["smiles"] == ADDUCT:
            ts = Path(st["files"]["ts"])
            shutil.copy(a.run_dir / ts.relative_to(ts.parts[0]), a.out / f"ts_{k}.xyz")
print(sorted(p.name for p in a.out.glob("*.xyz")))