"""bp_lm.py — the plain-BP standard-transformer reference with weight-EMA (task #34), rebuilt for
metric parity with the EP lineage: SAME parameter tensors as EQBlock (7.48M), SAME data + eval
protocol, and the SAME pema weight-EMA so both lineages report best-of(raw, ema) val. Forward = one
standard pre-LN residual block (no relaxation): h = x + attn(LN1 x) + ffn(LN2 x); logits = h @ Wh.
Historical anchor (no EMA, no qknorm): 'BP standard transformer (7.48M, lr 3e-3) best 1.7921 (20k)'
(FINDINGS:355). Run with --qknorm for architecture parity with the current EP recipe."""
import argparse, time, torch, torch.nn.functional as F
import lt_ep_train as L


def fwd(blk, idx):
    x = blk.embed(idx)
    h1 = F.layer_norm(x, (blk.C,), blk.ln1g, blk.ln1b)
    h2 = F.layer_norm(x, (blk.C,), blk.ln2g, blk.ln2b)
    h = x + blk.attn(h1) + (F.gelu(h2 @ blk.fc + blk.fcb, approximate='tanh') @ blk.pj + blk.pjb)
    return h @ (blk.tok.t() if getattr(blk, 'tie', False) else blk.Wh)


def evaluate(blk, nb=8, B=32):
    tot = 0.0
    with torch.no_grad():
        for _ in range(nb):
            idx, y = L.get_batch('val', B, blk.T)
            tot += float(F.cross_entropy(fwd(blk, idx).reshape(-1, L.vocab), y.reshape(-1)))
    return tot / nb


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--steps', type=int, default=32000)
    ap.add_argument('--B', type=int, default=24)
    ap.add_argument('--lr', type=float, default=3e-3)
    ap.add_argument('--wd', type=float, default=1e-4)
    ap.add_argument('--pema', type=float, default=0.999)
    ap.add_argument('--qknorm', action='store_true')
    ap.add_argument('--stdinit', action='store_true')          # standard transformer init (EQBlock's is tuned for relaxation)
    ap.add_argument('--beta2', type=float, default=0.999)
    ap.add_argument('--sched', choices=['cos', 'const'], default='cos')
    ap.add_argument('--tie', action='store_true')              # tok/Wh weight tying (standard small-LM trick)
    ap.add_argument('--log', type=int, default=200)
    ap.add_argument('--ckpt', type=str, default='runs/bp_lm.pt')
    cfg = ap.parse_args()
    torch.manual_seed(0)
    blk = L.EQBlock(512, 16, 256, 256, c=1.0, attn_mode='thick')
    blk.qknorm = cfg.qknorm
    if cfg.stdinit:                                            # GPT-style: N(0,0.02), scaled residual projections
        with torch.no_grad():
            for W in (blk.WQ, blk.WK, blk.WV, blk.fc, blk.Wh, blk.tok):
                W.normal_(0, 0.02)
            for W in (blk.WO, blk.pj):
                W.normal_(0, 0.02 / (2 ** 0.5))
            blk.pos.normal_(0, 0.01)
    if cfg.tie:                                                # tie head to embedding: Wh := tok^T, single parameter
        with torch.no_grad():
            blk.tok.copy_(0.5 * (blk.tok + blk.Wh.t()))
        blk.Wh = None                                          # fwd() will use tok.t() when tie is on
        blk.tie = True
        blk.allp = [p for p in blk.allp if p is not blk.Wh]
        blk.allp = blk.block + []                              # block already contains tok; Wh dropped
    opt = torch.optim.AdamW(blk.allp, lr=cfg.lr, weight_decay=cfg.wd, betas=(0.9, cfg.beta2))
    if cfg.sched == 'cos':
        sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, cfg.steps, eta_min=cfg.lr * 0.05)
    else:
        sched = torch.optim.lr_scheduler.LambdaLR(opt, lambda s: 1.0)
    ema = [p.detach().clone() for p in blk.allp]
    best, t0 = float('inf'), time.time()
    for step in range(1, cfg.steps + 1):
        idx, y = L.get_batch('train', cfg.B, blk.T)
        loss = F.cross_entropy(fwd(blk, idx).reshape(-1, L.vocab), y.reshape(-1))
        opt.zero_grad(set_to_none=True)
        loss.backward()
        opt.step(); sched.step()
        with torch.no_grad():
            for e, p in zip(ema, blk.allp):
                e.mul_(cfg.pema).add_(p, alpha=1 - cfg.pema)
        if step % cfg.log == 0:
            raw = evaluate(blk)
            bak = [p.detach().clone() for p in blk.allp]
            with torch.no_grad():
                for p, e in zip(blk.allp, ema):
                    p.copy_(e)
            emav = evaluate(blk)
            with torch.no_grad():
                for p, b in zip(blk.allp, bak):
                    p.copy_(b)
            m = min(raw, emav)
            if m < best:
                best = m
                torch.save({'allp': [p.detach().cpu() for p in blk.allp], 'step': step, 'best': best}, cfg.ckpt)
            print(f"step {step:5d}/{cfg.steps} | val {raw:.4f} ema {emav:.4f} (best {best:.4f}) "
                  f"| {step / (time.time() - t0):.2f} it/s", flush=True)
    print(f"[bp_lm] DONE best {best:.4f}", flush=True)


if __name__ == '__main__':
    main()