"""Retro + live wandb sync from trainer logs (runs launched before wandb default-on).
History logs replayed once; fw72m.log followed until DONE. Multi-run files (nsize_ladder)
switch tag on '[tag] cascade-EP' headers. Project: ept-cascade."""
import re, time
from pathlib import Path
import wandb
ENTITY = 'eqprop-llm-training'

R = Path('/home/yurenh2/ept/ep_run/runs')
EP_RE = re.compile(
    r'^step\s+(\d+)/\d+ \| train ([\d.]+) val ([\d.]+) \(best ([\d.]+)\) '
    r'\| beta=([\d.e+-]+) K=(\d+) skips=(\d+)\(d(\d+)/g(\d+)/r(\d+)\)(?: cos=([\d.]+))? '
    r'drift=([\d.]+) gn=([\d.e+-]+) sig=([\d.]+) \| ([\d.]+) it/s')
BP_RE = re.compile(r'^step\s+(\d+)/\d+ \| train ([\d.]+) val ([\d.]+) \(best ([\d.]+)\) \| ([\d.]+) it/s')
HDR = re.compile(r'^\[(\w+)\] cascade-')

def parse(line):
    m = EP_RE.match(line)
    if m:
        g = m.groups()
        d = {'train_ce': float(g[1]), 'val_ce': float(g[2]), 'best': float(g[3]),
             'beta_t': float(g[4]), 'K': int(g[5]), 'skips': int(g[6]),
             'drift': float(g[11]), 'gn': float(g[12]), 'sig': float(g[13]), 'it_s': float(g[14])}
        if g[10]: d['gate_cos'] = float(g[10])
        return int(g[0]), d
    m = BP_RE.match(line)
    if m:
        g = m.groups()
        return int(g[0]), {'train_ce': float(g[1]), 'val_ce': float(g[2]), 'best': float(g[3]), 'it_s': float(g[4])}
    return None

def sync_file(path, default_tag, follow=False):
    tag, run, last = default_tag, None, -1
    def ensure(t):
        nonlocal run, last
        if run is not None: run.finish()
        run = wandb.init(entity=ENTITY, project='ept-cascade', name=t, id=t, resume='allow', reinit=True)
        last = -1
    f = open(path, errors='replace')
    ensure(tag)
    while True:
        line = f.readline()
        if not line:
            if not follow: break
            if 'DONE best' in open(path, errors='replace').read()[-2000:]: break
            time.sleep(60); continue
        h = HDR.match(line)
        if h and h.group(1) != tag:
            tag = h.group(1); ensure(tag)
        p = parse(line)
        if p and p[0] > last:
            run.log(p[1], step=p[0]); last = p[0]
    if run is not None: run.finish()
    print(f'[sync] {path} done', flush=True)

HISTORY = [
    ('stage1b_ep_muon.log', 'stage1b_ep_muon'), ('stage1b_amp.log', 'stage1b_amp'),
    ('fw_smoke.log', 'fw_smoke'), ('amp_s1.log', 'amp_s1'), ('amp_s2.log', 'amp_s2'),
    ('amp_s3.log', 'amp_s3'), ('nsize_ladder.log', 'nsize_c256'),
    ('bsign_s1.log', 'bsign_s1'), ('bsign_s2.log', 'bsign_s2'), ('bsign_s3.log', 'bsign_s3'),
    ('f3e3cont.log', 'f3e3cont'),
    ('fw72m.log', 'fw72m'), ('fw72m_bp.log', 'fw72m_bp'), ('fw72m_f3e3.log', 'fw72m_f3e3'),
    ('fw72m_c.log', 'fw72m_c'), ('fw72m_c2.log', 'fw72m_c2'),
    ('w2_ctl.log', 'w2_ctl'), ('w2_blow.log', 'w2_blow'), ('w2_geta07.log', 'w2_geta07'),
    ('w2_geta05_k5.log', 'w2_geta05_k5'), ('w2_bcap.log', 'w2_bcap'), ('w2_adapt.log', 'w2_adapt'),
    ('arm_ctl.log', 'arm_ctl'), ('arm_mom99_f3e4.log', 'arm_mom99_f3e4'),
    ('arm_mom995_f1e4.log', 'arm_mom995_f1e4'), ('arm_cent_f3e3.log', 'arm_cent_f3e3'),
    ('arm_rich_f1e3.log', 'arm_rich_f1e3'), ('arm_mom99_f1e3.log', 'arm_mom99_f1e3'),
    ('arm_adamw_f1e3.log', 'arm_adamw_f1e3'),
]
for fn, tag in HISTORY:
    p = R / fn
    if p.exists():
        try: sync_file(p, tag)
        except Exception as e: print(f'[sync] {fn} FAILED: {e}', flush=True)
print('[sync] ALL DONE', flush=True)