import time, os, re, shutil
os.chdir("/home/yurenh2/ept/ep_run")
LOG, CK = "runs/ep_redx.log", "runs/ep_redx.pt"
os.makedirs("runs/redx_traj", exist_ok=True)
seen=set(); t0=time.time()
while time.time()-t0 < 8*3600:
    time.sleep(20)
    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+)/.*res=([\d.eE+-]+)", ls[-1])
    if not m: continue
    step=int(m.group(1)); res=float(m.group(2))
    if step not in seen and os.path.exists(CK) and os.path.getsize(CK)>1e6:
        try: shutil.copy2(CK, f"runs/redx_traj/s{step}.pt"); seen.add(step); print(f"froze s{step} res={res:.2e}", flush=True)
        except Exception: pass
    if res>0.25:
        print(f"DIVERGED step {step} res={res:.2e}; frozen={sorted(seen)}", flush=True); break
print("freezer done; frozen:", sorted(seen))