"""Backfill historical run logs into W&B (one wandb run per log, replayed at original steps).
Handles both trainer formats:
  ep:  step  100/32000 | val CE 2.99 ema=3.29 (best 2.99) | jr=0.1 res=8.7e-03 [rho=0.85] | 0.09 it/s
  bp:  step   200/32000 | val 3.4503 ema 7.9637 (best 3.4503) | 16.47 it/s
Governor / ABORT / DONE lines land in run.summary. Re-runnable: id = run name (resume overwrite).
Usage: python3 backfill_wandb.py            # backfill the curated map below
"""
import re, glob, os
import wandb

EP = re.compile(r'step\s+(\d+)/\d+ \| val CE ([\d.]+)(?: ema=([\d.]+))? \(best ([\d.]+)\) \| jr=([\d.]+) res=([\d.eE+-]+)(?: rho=([\d.]+))?.* ([\d.]+) it/s')
BP = re.compile(r'step\s+(\d+)/\d+ \| val ([\d.]+)(?: ema ([\d.]+))? \(best ([\d.]+)\).* ([\d.]+) it/s')

MAP = {  # glob -> (project, strip_suffixes)
    'runs/gov_s1[1-5].log':        'ept-c512',
    'runs/abl_*.log':              'ept-c512',
    'runs/ep_warm_fast.log':       'ept-c512',
    'runs/self_restart*.log':      'ept-c512',
    'runs/bp_lm.log':              'ept-c512',
    'runs/tol2_*.log':             'ept-tol',
    'runs/tol_*.log':              'ept-tol',
    'runs/par_*.log':              'ept-c512',
    'runs/delta_mirror/rung33m*.out': 'ept-33m',
}

def backfill(path, project):
    name = re.sub(r'_\d{6,}$', '', os.path.splitext(os.path.basename(path))[0])
    rows, events, done = [], [], None
    for ln in open(path, errors='replace'):
        m = EP.search(ln)
        if m:
            s, v, e, b, jr, res, rho, ips = m.groups()
            r = {'val_ce': float(v), 'best': float(b), 'jr': float(jr), 'res': float(res), 'it_per_s': float(ips)}
            if e: r['ema_ce'] = float(e)
            if rho: r['rho'] = float(rho)
            rows.append((int(s), r)); continue
        m = BP.search(ln)
        if m:
            s, v, e, b, ips = m.groups()
            r = {'val_ce': float(v), 'best': float(b), 'it_per_s': float(ips)}
            if e: r['ema_ce'] = float(e)
            rows.append((int(s), r)); continue
        if '[governor]' in ln or 'ABORT' in ln:
            events.append(ln.strip())
        if 'DONE best val CE' in ln:
            done = float(re.search(r'DONE best val CE ([\d.]+)', ln).group(1))
    if not rows:
        print(f'  skip (no rows): {path}'); return
    rows.sort(key=lambda t: t[0])
    last = -1
    run = wandb.init(project=project, name=name, id=name, resume='allow',
                     tags=['backfill'], config={'source_log': path}, reinit=True)
    for s, r in rows:
        if s <= last: continue   # resumed-run overlap: keep first pass, monotone steps only
        run.log(r, step=s); last = s
    bestv = done if done is not None else min(r['best'] for _, r in rows)
    run.summary['best_val_ce'] = bestv
    run.summary['final_step'] = last
    if events: run.summary['events'] = events[:40]
    if done is not None: run.summary['done'] = True
    run.finish()
    print(f'  {project}/{name}: {len(rows)} pts, best {bestv}, events {len(events)}')

seen = set()
for pat, project in MAP.items():
    for p in sorted(glob.glob(pat)):
        if p in seen: continue
        seen.add(p)
        try: backfill(p, project)
        except Exception as e: print(f'  ERROR {p}: {e}')
print(f'backfill complete: {len(seen)} logs')