#!/usr/bin/env python3
"""
Experiment 5: Relaxation Mechanism Ablation

Falsifies or confirms whether weight relaxation is structural prediction (BVP)
or generic regularization. Five seeds throughout.

Configs (all d=1024, 4L, 1000 steps, 10% active):
  1. dense
  2. dense + relax_graph (alpha=0.1)
  3. dense + relax_roll (alpha=0.1)
  4. ema_only
  5. ema + relax_graph (alpha=0.1)
  6. ema + relax_roll (alpha=0.1)
  7. ema + relax_random (random similarity matrix)
  8. ema + relax_shuffled_graph (real stats, broken structure)

Alpha sweep (graph, ema): 0.0, 0.01, 0.05, 0.1, 0.2, 0.5, 0.9

Diagnostics per run:
  - Val loss every 50 steps
  - grad_cos: cosine similarity between relaxer delta and dense oracle gradient
  - mag_ratio: ||relaxer_delta|| / ||oracle_grad|| on inactive chunks
  - mask_jaccard: step-to-step overlap of active set
"""
import argparse,json,math,os,random,sys,time,urllib.request
from collections import defaultdict
import torch,torch.nn as nn,torch.nn.functional as F
import tiktoken
print("imports ok",flush=True)

class Corpus:
    _i=None
    @classmethod
    def get(cls,bs,dev):
        if cls._i is None: cls._i=cls(bs,dev)
        return cls._i
    def __init__(self,bs,dev):
        self.block_size,self.device=bs,dev
        p="input.txt"
        if not os.path.exists(p):
            urllib.request.urlretrieve("https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt",p)
        enc=tiktoken.get_encoding("gpt2"); t=enc.encode(open(p).read())
        self.vocab_size=enc.n_vocab; d=torch.tensor(t,dtype=torch.long)
        si=int(0.9*len(d)); self.train_data,self.val_data=d[:si],d[si:]
        print(f"Corpus: V={self.vocab_size} train={len(self.train_data):,} val={len(self.val_data):,}",flush=True)
    def get_batch(self,split,bs,gen=None):
        d=self.train_data if split=="train" else self.val_data
        ix=torch.randint(len(d)-self.block_size-1,(bs,),generator=gen)
        x=torch.stack([d[i:i+self.block_size] for i in ix])
        y=torch.stack([d[i+1:i+self.block_size+1] for i in ix])
        return x.to(self.device),y.to(self.device)
def mg(s):
    g=torch.Generator(device="cpu"); g.manual_seed(s); return g

class SparseBwd(torch.autograd.Function):
    @staticmethod
    def forward(ctx,x,w,b,ac,cs,sdx):
        ctx.save_for_backward(x,w,ac); ctx.hb=b is not None; ctx.sdx=sdx; ctx.cs=cs
        return F.linear(x,w,b)
    @staticmethod
    def backward(ctx,gy):
        x,w,ac=ctx.saved_tensors; cs=ctx.cs
        xf=x.reshape(-1,x.shape[-1]); gf=gy.reshape(-1,gy.shape[-1])
        gw=torch.zeros_like(w)
        gb=torch.zeros(w.shape[0],device=w.device,dtype=w.dtype) if ctx.hb else None
        gx=torch.zeros_like(xf) if ctx.sdx else gf@w
        for c in ac.tolist():
            s,e=c*cs,(c+1)*cs; sl=gf[:,s:e]
            gw[s:e]=sl.t()@xf
            if gb is not None: gb[s:e]=sl.sum(0)
            if ctx.sdx: gx+=sl@w[s:e]
        return gx.reshape(x.shape),gw,gb,None,None,None

class SL(nn.Linear):
    def __init__(self,i,o,bias=True):
        super().__init__(i,o,bias=bias)
        self.se=False; self.sdx=False; self.ac=None; self.cs=64
    def forward(self,x):
        if not self.se or self.ac is None: return F.linear(x,self.weight,self.bias)
        return SparseBwd.apply(x,self.weight,self.bias,self.ac,self.cs,self.sdx)

class Attn(nn.Module):
    def __init__(self,d,nh,bs,do):
        super().__init__(); self.nh=nh; self.hd=d//nh
        self.qkv=SL(d,3*d); self.proj=SL(d,d); self.drop=nn.Dropout(do)
        self.register_buffer("mask",torch.tril(torch.ones(bs,bs)).view(1,1,bs,bs))
    def forward(self,x):
        B,T,C=x.shape; q,k,v=self.qkv(x).split(C,2)
        q=q.view(B,T,self.nh,self.hd).transpose(1,2)
        k=k.view(B,T,self.nh,self.hd).transpose(1,2)
        v=v.view(B,T,self.nh,self.hd).transpose(1,2)
        a=(q@k.transpose(-2,-1))/math.sqrt(self.hd)
        a=a.masked_fill(self.mask[:,:,:T,:T]==0,float("-inf"))
        a=self.drop(F.softmax(a,dim=-1))
        return self.proj((a@v).transpose(1,2).contiguous().view(B,T,C))

class FFN(nn.Module):
    def __init__(self,d,do):
        super().__init__(); self.fc=SL(d,4*d); self.proj=SL(4*d,d); self.drop=nn.Dropout(do)
    def forward(self,x): return self.drop(self.proj(F.gelu(self.fc(x))))

class Blk(nn.Module):
    def __init__(self,d,nh,bs,do):
        super().__init__(); self.ln1=nn.LayerNorm(d); self.attn=Attn(d,nh,bs,do)
        self.ln2=nn.LayerNorm(d); self.mlp=FFN(d,do)
    def forward(self,x): x=x+self.attn(self.ln1(x)); return x+self.mlp(self.ln2(x))

class GPT(nn.Module):
    def __init__(self,V,bs,nl,nh,d,do):
        super().__init__(); self.te=nn.Embedding(V,d); self.pe=nn.Embedding(bs,d)
        self.blocks=nn.Sequential(*[Blk(d,nh,bs,do) for _ in range(nl)])
        self.ln=nn.LayerNorm(d); self.head=nn.Linear(d,V)
    def forward(self,idx,tgt=None):
        B,T=idx.shape; x=self.te(idx)+self.pe(torch.arange(T,device=idx.device))[None]
        lo=self.head(self.ln(self.blocks(x)))
        return lo,F.cross_entropy(lo.view(-1,lo.size(-1)),tgt.view(-1)) if tgt is not None else None
    def np(self): return sum(p.numel() for p in self.parameters())
def gsl(m): return [x for x in m.modules() if isinstance(x,SL)]

class Sched:
    def __init__(self,model,frac,cs,dev,beta=0.95,sim_hist=128,min_sim=8):
        self.frac,self.cs,self.dev,self.beta=frac,cs,dev,beta
        self.sim_hist,self.min_sim=sim_hist,min_sim
        self.lins=gsl(model); self.m2i,self.m2l={},{}; off=0
        for m in self.lins:
            m.cs=cs; nc=m.out_features//cs; assert m.out_features%cs==0
            self.m2i[m]=torch.arange(off,off+nc,device=dev)
            self.m2l[m]=torch.arange(nc,device=dev); off+=nc
        self.nc=off; self.ema=torch.zeros(self.nc,device=dev)
        self.act=torch.zeros(self.nc,dtype=torch.bool,device=dev)
        self.prev_act=torch.zeros(self.nc,dtype=torch.bool,device=dev)
        self.mass_history=[]; self.similarity=None
    def gf(self,step,wu,an):
        if step<wu: return 1.0
        if an>0 and step<wu+an:
            p=(step-wu)/an; return self.frac+(1-self.frac)*0.5*(1+math.cos(math.pi*p))
        return self.frac
    def choose(self,step,wu,an):
        self.prev_act=self.act.clone()
        f=self.gf(step,wu,an)
        if f>=0.999: self.act.fill_(True); self._inst(); return
        k=max(1,int(f*self.nc)); self.act.fill_(False)
        idx=torch.topk(self.ema+1e-9*torch.rand_like(self.ema),k=k).indices
        self.act[idx]=True; self._inst()
    def _inst(self):
        for m,gi in self.m2i.items(): m.ac=self.m2l[m][self.act[gi]]
    @torch.no_grad()
    def update(self,step,wu):
        cur=torch.zeros_like(self.ema)
        for m,ids in self.m2i.items():
            if m.weight.grad is None: continue
            s=m.weight.grad.square().view(len(ids),self.cs,-1).sum((1,2))
            if m.bias is not None and m.bias.grad is not None:
                s+=m.bias.grad.square().view(len(ids),self.cs).sum(1)
            cur[ids]=torch.sqrt(s+1e-30)
        obs=self.act; new=obs&(self.ema==0); old=obs&~new
        self.ema[new]=cur[new]; self.ema[old]=self.beta*self.ema[old]+(1-self.beta)*cur[old]
        if step<wu:
            self.mass_history.append(cur.clone())
            if len(self.mass_history)>self.sim_hist:
                self.mass_history=self.mass_history[-self.sim_hist:]
            if len(self.mass_history)>=self.min_sim:
                self._build_sim()
        return cur
    def _build_sim(self):
        H=torch.stack(self.mass_history)
        H=(H-H.mean(0,keepdim=True))/(H.std(0,keepdim=True)+1e-6)
        S=torch.clamp((H.T@H)/max(1,H.shape[0]-1),min=0)
        S.fill_diagonal_(0)
        ok=torch.zeros_like(S,dtype=torch.bool)
        for _,ids in self.m2i.items(): ok[ids[:,None],ids[None,:]]=True
        self.similarity=torch.where(ok,S,torch.zeros_like(S))
    def mask_jaccard(self):
        if self.prev_act.sum()==0: return 0.0
        i=(self.act&self.prev_act).sum().item()
        u=(self.act|self.prev_act).sum().item()
        return i/max(u,1)

class GraphRelaxer:
    def __init__(self, sched, alpha=0.1, iters=3):
        self.sched,self.alpha,self.iters=sched,alpha,iters
    @torch.no_grad()
    def relax(self):
        S=self.sched.similarity
        if S is None: return {}
        act=self.sched.act; deltas={}
        for m,ids in self.sched.m2i.items():
            nc=len(ids); cs=self.sched.cs; di=m.weight.shape[1]
            S_local=S[ids][:,ids]
            rs=S_local.sum(1,keepdim=True)+1e-12; S_n=S_local/rs
            la=act[ids]; li=~la
            if li.sum()==0: continue
            W=m.weight.data.view(nc,cs,di); W_before=W[li].clone()
            for _ in range(self.iters):
                Wf=W.reshape(nc,-1); Wa=(S_n@Wf).view(nc,cs,di)
                W[li]=(1-self.alpha)*W[li]+self.alpha*Wa[li]
            m.weight.data=W.view(m.out_features,di)
            deltas[m]=W[li]-W_before
        return deltas

class RollRelaxer:
    def __init__(self, sched, alpha=0.1, iters=3):
        self.sched,self.alpha,self.iters=sched,alpha,iters
    @torch.no_grad()
    def relax(self):
        act=self.sched.act; deltas={}
        for m,ids in self.sched.m2i.items():
            nc=len(ids); cs=self.sched.cs; di=m.weight.shape[1]
            la=act[ids]; li=~la
            if li.sum()==0: continue
            W=m.weight.data.view(nc,cs,di); W_before=W[li].clone()
            for _ in range(self.iters):
                Wp=torch.roll(W,1,dims=0); Wn=torch.roll(W,-1,dims=0)
                Wa=(Wp+Wn)/2.0
                W[li]=(1-self.alpha)*W[li]+self.alpha*Wa[li]
            m.weight.data=W.view(m.out_features,di)
            deltas[m]=W[li]-W_before
        return deltas

class RandomRelaxer:
    def __init__(self, sched, alpha=0.1, iters=3):
        self.sched,self.alpha,self.iters=sched,alpha,iters
        self._rand_sim=None
    def _get_rand_sim(self):
        if self._rand_sim is not None: return self._rand_sim
        S=self.sched.similarity
        if S is None: return None
        R=torch.rand_like(S)*S.abs().mean()
        R.fill_diagonal_(0)
        ok=torch.zeros_like(R,dtype=torch.bool)
        for _,ids in self.sched.m2i.items(): ok[ids[:,None],ids[None,:]]=True
        self._rand_sim=torch.where(ok,R,torch.zeros_like(R))
        return self._rand_sim
    @torch.no_grad()
    def relax(self):
        S=self._get_rand_sim()
        if S is None: return {}
        act=self.sched.act; deltas={}
        for m,ids in self.sched.m2i.items():
            nc=len(ids); cs=self.sched.cs; di=m.weight.shape[1]
            S_local=S[ids][:,ids]
            rs=S_local.sum(1,keepdim=True)+1e-12; S_n=S_local/rs
            la=act[ids]; li=~la
            if li.sum()==0: continue
            W=m.weight.data.view(nc,cs,di); W_before=W[li].clone()
            for _ in range(self.iters):
                Wf=W.reshape(nc,-1); Wa=(S_n@Wf).view(nc,cs,di)
                W[li]=(1-self.alpha)*W[li]+self.alpha*Wa[li]
            m.weight.data=W.view(m.out_features,di)
            deltas[m]=W[li]-W_before
        return deltas

class ShuffledGraphRelaxer:
    def __init__(self, sched, alpha=0.1, iters=3):
        self.sched,self.alpha,self.iters=sched,alpha,iters
        self._shuf_sim=None
    def _get_shuf_sim(self):
        if self._shuf_sim is not None: return self._shuf_sim
        S=self.sched.similarity
        if S is None: return None
        Ss=S.clone()
        for _,ids in self.sched.m2i.items():
            n=len(ids)
            block=Ss[ids][:,ids].clone()
            perm=torch.randperm(n,device=S.device)
            block=block[perm][:,perm]
            block.fill_diagonal_(0)
            Ss[ids[:,None],ids[None,:]]=block
        self._shuf_sim=Ss
        return self._shuf_sim
    @torch.no_grad()
    def relax(self):
        S=self._get_shuf_sim()
        if S is None: return {}
        act=self.sched.act; deltas={}
        for m,ids in self.sched.m2i.items():
            nc=len(ids); cs=self.sched.cs; di=m.weight.shape[1]
            S_local=S[ids][:,ids]
            rs=S_local.sum(1,keepdim=True)+1e-12; S_n=S_local/rs
            la=act[ids]; li=~la
            if li.sum()==0: continue
            W=m.weight.data.view(nc,cs,di); W_before=W[li].clone()
            for _ in range(self.iters):
                Wf=W.reshape(nc,-1); Wa=(S_n@Wf).view(nc,cs,di)
                W[li]=(1-self.alpha)*W[li]+self.alpha*Wa[li]
            m.weight.data=W.view(m.out_features,di)
            deltas[m]=W[li]-W_before
        return deltas

class NullRelaxer:
    def relax(self): return {}

class CAdam:
    def __init__(self,model,lr=3e-4,cs=64):
        self.model,self.lr,self.cs=model,lr,cs
        self.st={}; self.p2m={}
        for m in gsl(model):
            if m.weight is not None: self.p2m[m.weight]=m
            if m.bias is not None: self.p2m[m.bias]=m
    def zero_grad(self):
        for p in self.model.parameters(): p.grad=None
    @torch.no_grad()
    def step(self):
        for p in self.model.parameters():
            if p.grad is None: continue
            if p not in self.st: self.st[p]={"m":torch.zeros_like(p),"v":torch.zeros_like(p)}
            m,v=self.st[p]["m"],self.st[p]["v"]
            sm=self.p2m.get(p); ac=getattr(sm,'ac',None) if sm else None
            if ac is None:
                m.mul_(0.9).add_(p.grad,alpha=0.1); v.mul_(0.999).addcmul_(p.grad,p.grad,value=0.001)
                p.sub_(m/(torch.sqrt(v)+1e-8),alpha=self.lr)
            else:
                m.mul_(0.9).add_(p.grad,alpha=0.1); v.mul_(0.999).addcmul_(p.grad,p.grad,value=0.001)
                for c in ac.tolist():
                    s,e=c*self.cs,(c+1)*self.cs
                    p.data[s:e].sub_(m[s:e]/(torch.sqrt(v[s:e])+1e-8),alpha=self.lr)

@torch.no_grad()
def ev(model,corpus,bs,n=20,seed=9999):
    model.eval(); ls=[model(*corpus.get_batch("val",bs,mg(seed+i)))[1].item() for i in range(n)]
    model.train(); a=sum(ls)/len(ls); return a,math.exp(min(a,20))

def compute_relaxer_diagnostics(model, sched, relaxer_deltas, x, y, corpus, bs, cs):
    """
    Compare relaxer delta on inactive chunks to what dense gradient would have been.
    Returns (grad_cos, mag_ratio) or (None, None) if not applicable.
    """
    if not relaxer_deltas: return None, None

    # Need gradients enabled for the dense forward/backward
    for m in gsl(model): m.se=False
    for p in model.parameters(): p.grad=None
    with torch.enable_grad():
        _,lo=model(x,y)
        lo.backward()

    cos_sims=[]; mag_ratios=[]
    with torch.no_grad():
        for m,delta in relaxer_deltas.items():
            if m not in sched.m2i: continue
            ids=sched.m2i[m]; nc=len(ids); di=m.weight.shape[1]
            la=sched.act[ids]; li=~la
            if li.sum()==0 or m.weight.grad is None: continue
            dense_g=m.weight.grad.view(nc,cs,di)[li]
            d_flat=delta.reshape(-1); g_flat=dense_g.reshape(-1)
            dn=d_flat.norm(); gn=g_flat.norm()
            if dn>1e-12 and gn>1e-12:
                cos_sims.append(F.cosine_similarity(d_flat.unsqueeze(0),g_flat.unsqueeze(0)).item())
                mag_ratios.append((dn/gn).item())

    for m in gsl(model): m.se=True
    for p in model.parameters(): p.grad=None

    if not cos_sims: return None, None
    return sum(cos_sims)/len(cos_sims), sum(mag_ratios)/len(mag_ratios)

def run1(mode, steps, bs, bsz, nl, nh, d, cs, af, wu, an, lr, dev, seed,
         alpha=0.1, iters=3, diag_interval=100):
    torch.manual_seed(seed); random.seed(seed)
    if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)
    corpus=Corpus.get(bsz,dev)
    model=GPT(corpus.vocab_size,bsz,nl,nh,d,0.1).to(dev)
    for m in gsl(model): m.cs=cs

    is_dense = mode.startswith("dense")
    is_sparse = not is_dense
    needs_relax = "relax" in mode

    sched=None
    if is_sparse:
        sched=Sched(model,af,cs,dev)
    elif needs_relax:
        sched=Sched(model,af,cs,dev)

    opt=CAdam(model,lr,cs)

    if not needs_relax:
        relaxer=NullRelaxer()
    elif "random" in mode:
        relaxer=RandomRelaxer(sched,alpha,iters)
    elif "shuffled" in mode:
        relaxer=ShuffledGraphRelaxer(sched,alpha,iters)
    elif "roll" in mode:
        relaxer=RollRelaxer(sched,alpha,iters)
    elif "graph" in mode:
        relaxer=GraphRelaxer(sched,alpha,iters)
    else:
        relaxer=NullRelaxer()

    np_=model.np()
    val_curve=[]; grad_cos_log=[]; mag_ratio_log=[]; jaccard_log=[]

    if dev=="cuda": torch.cuda.synchronize()
    t0=time.perf_counter()

    for step in range(steps):
        x,y=corpus.get_batch("train",bs,mg(step))

        if is_sparse:
            sched.choose(step,wu,an)
            for m in gsl(model): m.se=True; m.sdx=False
        else:
            for m in gsl(model): m.se=False; m.ac=None
            if sched:
                sched.choose(step,wu,an)

        opt.zero_grad(); _,loss=model(x,y); loss.backward()

        if sched:
            sched.update(step,wu)
            jaccard_log.append((step, sched.mask_jaccard()))

        opt.step()

        relax_deltas={}
        if needs_relax and step>=wu+an:
            if is_dense and sched:
                for m,ids in sched.m2i.items():
                    m.ac=sched.m2l[m][sched.act[ids]]
            relax_deltas=relaxer.relax()
            if is_dense and sched:
                for m in gsl(model): m.ac=None

        if step%50==0:
            vl,_=ev(model,corpus,bs,n=10,seed=7777)
            val_curve.append((step,vl))

        if step%diag_interval==0 and step>=wu+an and relax_deltas and sched:
            gc,mr=compute_relaxer_diagnostics(model,sched,relax_deltas,x,y,corpus,bs,cs)
            if gc is not None:
                grad_cos_log.append((step,gc))
                mag_ratio_log.append((step,mr))

        if step%200==0:
            print(f"    step {step}/{steps} loss={loss.item():.4f}",flush=True)

    if dev=="cuda": torch.cuda.synchronize()
    wall=time.perf_counter()-t0
    for m in gsl(model): m.se=False
    vl,vp=ev(model,corpus,bs,n=30)
    del model; torch.cuda.empty_cache() if dev=="cuda" else None

    return {
        "vl":vl,"vp":vp,"wall":wall,"ms":1000*wall/steps,"np":np_,"tl":loss.item(),
        "val_curve":val_curve,"grad_cos":grad_cos_log,"mag_ratio":mag_ratio_log,
        "jaccard":jaccard_log,
    }

def runs(cfg,seeds):
    rs=[]
    for s in seeds:
        c=dict(cfg); c["seed"]=s; rs.append(run1(**c))
    vls=[r["vl"] for r in rs]; ml=sum(vls)/len(vls)
    sl=(sum((x-ml)**2 for x in vls)/max(1,len(vls)-1))**0.5
    return {"ml":ml,"sl":sl,"rs":rs,"ms":sum(r["ms"] for r in rs)/len(rs)}

def main():
    p=argparse.ArgumentParser()
    p.add_argument("--device",default="cuda"); p.add_argument("--steps",type=int,default=1000)
    p.add_argument("--seeds",default="42,123,456,789,1024")
    p.add_argument("--d",type=int,default=1024); p.add_argument("--nl",type=int,default=4)
    p.add_argument("--nh",type=int,default=8); p.add_argument("--bs",type=int,default=8)
    p.add_argument("--bsz",type=int,default=256); p.add_argument("--cs",type=int,default=64)
    p.add_argument("--af",type=float,default=0.10); p.add_argument("--wu",type=int,default=50)
    p.add_argument("--an",type=int,default=200); p.add_argument("--lr",type=float,default=3e-4)
    a=p.parse_args(); seeds=[int(s) for s in a.seeds.split(",")]

    if a.device=="cuda" and torch.cuda.is_available():
        print(f"GPU: {torch.cuda.get_device_name()} VRAM: {torch.cuda.get_device_properties(0).total_memory/1e9:.1f}GB",flush=True)
    print(f"d={a.d} nl={a.nl} steps={a.steps} seeds={seeds}",flush=True)
    print(f"cs={a.cs} af={a.af} wu={a.wu} an={a.an} lr={a.lr}",flush=True)

    base=dict(steps=a.steps,bs=a.bs,bsz=a.bsz,nl=a.nl,nh=a.nh,d=a.d,cs=a.cs,af=a.af,
              wu=a.wu,an=a.an,lr=a.lr,dev=a.device,alpha=0.1,iters=3)

    configs=[
        ("dense",                "dense"),
        ("dense+relax_graph",    "dense+relax_graph"),
        ("dense+relax_roll",     "dense+relax_roll"),
        ("ema_only",             "ema_only"),
        ("ema+relax_graph",      "ema+relax_graph"),
        ("ema+relax_roll",       "ema+relax_roll"),
        ("ema+relax_random",     "ema+relax_random"),
        ("ema+relax_shuffled",   "ema+relax_shuffled_graph"),
    ]

    print("\n"+"="*80,flush=True)
    print("EXP 5: Relaxation Mechanism Ablation (5 seeds)",flush=True)
    print("="*80,flush=True)

    R={}
    for name,mode in configs:
        print(f"\n--- {name} ({len(seeds)} seeds) ---",flush=True)
        R[name]=runs({**base,"mode":mode},seeds)

    print(f"\n{'Method':<25} | {'Val Loss':>20} | {'ms/step':>8}",flush=True)
    print("-"*60,flush=True)
    for name,_ in configs:
        r=R[name]
        print(f"{name:<25} | {r['ml']:.4f} ± {r['sl']:.4f}      | {r['ms']:>7.1f}",flush=True)

    print(f"\n--- Alpha sweep (ema+relax_graph, 5 seeds) ---",flush=True)
    print(f"{'alpha':>6} | {'Val Loss':>20} | {'ms/step':>8}",flush=True)
    print("-"*42,flush=True)
    alpha_results={}
    for alpha in [0.0, 0.01, 0.05, 0.1, 0.2, 0.5, 0.9]:
        r=runs({**base,"mode":"ema+relax_graph","alpha":alpha},seeds)
        alpha_results[alpha]=r
        print(f"{alpha:>6.2f} | {r['ml']:.4f} ± {r['sl']:.4f}      | {r['ms']:>7.1f}",flush=True)

    print(f"\n--- Diagnostics (grad_cos, mag_ratio) ---",flush=True)
    for name in ["ema+relax_graph","ema+relax_roll","ema+relax_random","ema+relax_shuffled"]:
        if name not in R: continue
        gc_all=[]; mr_all=[]
        for res in R[name]["rs"]:
            gc_all.extend([x[1] for x in res["grad_cos"]])
            mr_all.extend([x[1] for x in res["mag_ratio"]])
        if gc_all:
            gc_m=sum(gc_all)/len(gc_all); mr_m=sum(mr_all)/len(mr_all)
            print(f"  {name:<25}: grad_cos={gc_m:.4f}  mag_ratio={mr_m:.4f}",flush=True)

    all_results={"configs":R,"alpha_sweep":alpha_results}
    with open("exp5.json","w") as f:
        json.dump(all_results,f,indent=2,default=str)
    print("\n✓ exp5.json saved",flush=True)
    print(f"\nTotal: {len(configs)*len(seeds) + 7*len(seeds)} runs",flush=True)

if __name__=="__main__": main()