"""Consolidated trajectory freezer for the 4-seed dip farm (runs anywhere on the shared NFS):
copies runs/dipfarm_sN.pt -> runs/dipfarm_sN_traj/s<step>.pt whenever a new step line appears.
Feeds the ARPACK dip-screening (stability-dip statistics: rate, width, screenable yield)."""
import time, os, re, shutil

os.chdir("/home/yurenh2/ept/ep_run")
SEEDS = (1, 2, 3, 4)
seen = {s: set() for s in SEEDS}
for s in SEEDS:
    os.makedirs(f"runs/dipfarm_s{s}_traj", exist_ok=True)
t0 = time.time()
while time.time() - t0 < 24 * 3600:
    time.sleep(20)
    done = 0
    for s in SEEDS:
        log, ck = f"runs/dipfarm_s{s}.log", f"runs/dipfarm_s{s}.pt"
        try:
            ls = [l for l in open(log) if l.startswith("step")]
        except Exception:
            continue
        if not ls:
            continue
        m = re.search(r"step +(\d+)/", ls[-1])
        if not m:
            continue
        step = int(m.group(1))
        if step not in seen[s] and os.path.exists(ck) and os.path.getsize(ck) > 1e6:
            try:
                shutil.copy2(ck, f"runs/dipfarm_s{s}_traj/s{step}.pt")
                seen[s].add(step)
                print(f"froze s{s}/{step}", flush=True)
            except Exception:
                pass
        if step >= 4000 or "DONE" in ls[-1]:
            done += 1
    if done == len(SEEDS):
        break
print("dipfarm freezer done:", {s: len(v) for s, v in seen.items()}, flush=True)