"""Combine the two Geodesic sweeps into geodesic_combined_scores.json.

    python combine_geodesic.py      (run in this folder, after scoring both sweeps)

Reads geodesic_mapped_scores.json and geodesic_unmapped_scores.json (score_channels output for
the searches with and without atom mapping) and the run folders geodesic_runs_mapped/ and
geodesic_runs_unmapped/. A reaction's routes are those of both sweeps. Routes with an IRC end
where a hydrogen sits within 1.45 A of two heavy atoms are dropped, because bond perception
misreads bridging hydrogens. The best route is the lowest remaining one.

This step was first run as an inline script on 2026-10-04; this file is that code, with
file and folder names made neutral.
"""
import json, numpy as np
from pathlib import Path
def frames(p):
    L=p.read_text().splitlines(); out=[]; i=0
    while i<len(L):
        if not L[i].strip(): i+=1; continue
        n=int(L[i].split()[0]); blk=[l.split() for l in L[i+2:i+2+n]]
        out.append(([b[0] for b in blk], np.array([[float(v) for v in b[1:4]] for b in blk]))); i+=n+2
    return out
def bridged(s,x):
    hv=[j for j,e in enumerate(s) if e!="H"]
    return any(len(d:=sorted(np.linalg.norm(x[hv]-x[h],axis=1)))>1 and d[1]<1.45 for h in [i for i,e in enumerate(s) if e=="H"])
def bad_groups(root, rid):
    run=next(Path(root).glob(f"rxn_*_{rid}"),None); bad={}
    if run:
        for p in list(run.glob("channels/*/irc.xyz"))+list(run.glob("alternate-channels/*/step_*/irc.xyz")):
            f=frames(p)
            if f and (bridged(*f[0]) or bridged(*f[-1])):
                g=p.parent.name if p.parent.name.startswith("channel") else p.parent.parent.name
                bad[g]="an IRC end has a bridging hydrogen (H within 1.45 A of two heavy atoms); bond perception misreads it"
    return bad
M={r["reaction_id"]:r for r in json.load(open("geodesic_mapped_scores.json"))}
N={r["reaction_id"]:r for r in json.load(open("geodesic_unmapped_scores.json"))}
out=[]
for rid in sorted(M):
    m,n=M[rid],N[rid]; rec=dict(m); excl=[]
    found=[]
    for root,r in (("geodesic_runs_mapped",m),("geodesic_runs_unmapped",n)):
        bad=bad_groups(root,rid)
        for f in r.get("found") or []:
            if f["group"] in bad: excl.append({"root":root,"group":f["group"],"why":bad[f["group"]]}); continue
            found.append({**f,"root":root})
    rec["found"]=found; rec["excluded"]=excl; rec["n_found"]=len(found)
    rec["returncode"]=0 if 0 in (m.get("returncode"),n.get("returncode")) else m.get("returncode")
    rec["wall_seconds"]=max(m.get("wall_seconds") or 0, n.get("wall_seconds") or 0)
    for k in ("best_barrier_kcal","best_kind","improvement_kcal","best_root"): rec.pop(k,None)
    if found:
        b=min(found,key=lambda f:f["barrier_kcal"])
        rec.update(best_barrier_kcal=b["barrier_kcal"],best_kind=b["kind"],best_root=b["root"],improvement_kcal=round(rec["ref_barrier_kcal"]-b["barrier_kcal"],3))
        low=lambda root: [f for f in found if f["root"]==root and rec["ref_barrier_kcal"]-f["barrier_kcal"]>1]
        rec["lower_via"]=[x for x,root in (("mapped","geodesic_runs_mapped"),("no mapping","geodesic_runs_unmapped")) if low(root)]
    out.append(rec)
json.dump(out,open("geodesic_combined_scores.json","w"),indent=1)
print(len(out), sum(1 for r in out if r.get("improvement_kcal",0)>1), [ (r["reaction_id"],r["lower_via"]) for r in out if r.get("improvement_kcal",0)>1])