"""How the 30 reactions were picked (run 2026-10-13 from the playgroud study repo:
`cd playgroud; python select_set30.py OUTDIR`). Per source set: keep reactions whose start and end
differ by 1-4 bonds (distance < 1.25 x summed covalent radii), sort by atom count, take 15 evenly
spaced in that order."""
import csv, sys, numpy as np
base = "data/reaction_qm/"
R = {"H":0.31,"C":0.76,"N":0.71,"O":0.66,"F":0.57,"P":1.07,"S":1.05,"Cl":1.02,"Si":1.11,"Br":1.2,"B":0.84,"Li":1.28,"Na":1.66,"Mg":1.41}
def read(p):
    l = open(p).read().splitlines(); n = int(l[0]); s=[x.split()[0] for x in l[2:2+n]]
    return s, np.array([[float(v) for v in x.split()[1:4]] for x in l[2:2+n]])
def nchange(r):
    s1, x1 = read(base + r["start_xyz"]); s2, x2 = read(base + r["end_xyz"])
    if s1 != s2: return 99
    def bonds(s, x):
        d = np.linalg.norm(x[:,None]-x[None], axis=-1); rr = np.array([R.get(a,1.0) for a in s])
        return {(i,j) for i in range(len(s)) for j in range(i+1,len(s)) if d[i,j] < 1.25*(rr[i]+rr[j])}
    return len(bonds(s1,x1) ^ bonds(s2,x2))
out = []
for f, tag in [("geodesic_gxtb.csv", "g"), ("kaist_irc_500.csv", "k")]:
    rows = [r for r in csv.DictReader(open("src/playgroud/data/" + f)) if 1 <= nchange(r) <= 4]
    rows.sort(key=lambda r: int(r["n_atoms"]))
    pick = [rows[int(i)] for i in np.linspace(0, len(rows) - 1, 15)]
    for r in pick:
        out.append({"idx": tag + r["idx"], "reaction_id": r["reaction_id"], "n_atoms": r["n_atoms"], "charge": r["charge"],
                    "multiplicity": r["multiplicity"], "barrier_kcal": r["barrier_kcal"], "start_xyz": r["start_xyz"],
                    "end_xyz": r["end_xyz"], "ts_reference_xyz": r["ts_reference_xyz"], "bond_changes": nchange(r)})
with open(sys.argv[1] + "/set30.csv", "w", newline="") as fh:
    w = csv.DictWriter(fh, fieldnames=list(out[0])); w.writeheader(); w.writerows(out)
print(len(out), "reactions; atoms:", [int(r["n_atoms"]) for r in out])