"""dp_ep.py — data-parallel EP (task #14's last item). EP has no autograd backward, so DDP does not
apply; this is the minimal manual-allreduce loop: identical init on every rank (same seed), per-rank
batches, ep_step -> all_reduce(mean) over the canonical param order (zeros for absent grads so the
collective stays aligned), all_reduce the residual so the jr controller stays in lockstep, identical
AdamW steps -> weights never diverge. Launch:
  torchrun --nproc_per_node=N dp_ep.py --steps 300 [recipe flags...]
Effective batch = B * world (lr scaling deliberately NOT applied for the smoke; tune later)."""
import argparse, os, time, torch
import torch.distributed as dist
import lt_ep_train as L


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--steps', type=int, default=300)
    ap.add_argument('--B', type=int, default=24)
    ap.add_argument('--lr', type=float, default=6e-4)
    ap.add_argument('--seed', type=int, default=0)
    ap.add_argument('--reg_delay', type=int, default=10 ** 9)  # smoke default: reg-free (Pascal-safe)
    ap.add_argument('--resreg', type=float, default=0.2)
    ap.add_argument('--jacreg', type=float, default=0.1)
    ap.add_argument('--t2sel', type=int, default=40)
    ap.add_argument('--hr', type=float, default=0.02)
    ap.add_argument('--holofast', action='store_true')
    ap.add_argument('--sdpa', action='store_true')
    ap.add_argument('--log', type=int, default=50)
    cfg = ap.parse_args()

    dist.init_process_group('nccl')
    rank, world = dist.get_rank(), dist.get_world_size()
    torch.cuda.set_device(rank)
    L.dev = f'cuda:{rank}'

    torch.manual_seed(cfg.seed)                            # identical init on every rank
    blk = L.EQBlock(512, 16, 256, 256, c=1.0, attn_mode='thick')
    blk.qknorm = True; blk.track = True
    blk.holofast, blk.sdpa = cfg.holofast, cfg.sdpa
    for p in blk.allp:
        dist.broadcast(p.data, 0)                          # belt & suspenders: exact bitwise identical start
    torch.manual_seed(cfg.seed * 1009 + rank + 1)          # per-rank data stream
    opt = torch.optim.AdamW(blk.allp, lr=cfg.lr, weight_decay=1e-4)

    jr = cfg.jacreg
    t0 = time.time()
    for step in range(1, cfg.steps + 1):
        idx, y = L.get_batch('train', cfg.B, blk.T)
        dly = step < cfg.reg_delay
        grads, res = L.ep_step(blk, idx, y, 150, 20, 0.1, 0.02, 0.0 if dly else jr, holo=2, hr=cfg.hr,
                               t1max=300, res_est=1e-4, t2sel=cfg.t2sel, corr_every=1, res_gate=0.0,
                               resreg=0.0 if dly else cfg.resreg)
        rt = torch.tensor([res], device=L.dev)
        dist.all_reduce(rt); res = float(rt) / world       # controller signal identical across ranks
        for p in blk.allp:                                 # canonical order: aligned collectives
            g = grads.get(id(p))
            if g is None:
                g = torch.zeros_like(p)
            dist.all_reduce(g)
            g /= world
            p.grad = g
        opt.step(); opt.zero_grad(set_to_none=True)
        if step % cfg.log == 0 and rank == 0:
            v = L.evaluate(blk, 150, 0.1, nb=2)
            print(f"[dp{world}] step {step}/{cfg.steps} | val {v:.4f} | res {res:.1e} | "
                  f"{step / (time.time() - t0):.3f} it/s(x{world}B)", flush=True)
    if rank == 0:
        print(f"[dp{world}] DONE {cfg.steps} steps in {time.time() - t0:.0f}s", flush=True)
    dist.destroy_process_group()


if __name__ == '__main__':
    main()