"""update_size_b2.py (run 5): like update_size.py but with --beta2, and logging Adam's internal scales:
rms_grad (per-coordinate RMS gradient), rms_mhat, rms_sqrt_vhat (RMS of bias-corrected sqrt(v)), rms_step_over_lr (actual RMS step / lr),
plus frac_sat. Shows whether the step is damped by eps (sqrt(v)~g << eps) or by a stale v (sqrt(v) >> g).
Usage: python update_size_b2.py --beta2 0.999 --norm 82 --eps 1e-8 --steps 4000 > update_size_b2-0.999_N82_eps1e-8.csv"""
import argparse, torch, math
torch.set_num_threads(1)
import torch.nn as nn, torch.nn.functional as F
P=97; D_EMB,D_HID=64,256; FRAC=0.5
ap=argparse.ArgumentParser(); ap.add_argument('--norm',type=float,default=88); ap.add_argument('--seed',type=int,default=0); ap.add_argument('--steps',type=int,default=3000); ap.add_argument('--lr',type=float,default=1e-3); ap.add_argument('--eps',type=float,default=1e-8); ap.add_argument('--beta2',type=float,default=0.98); args=ap.parse_args()
BETAS=(0.9,args.beta2)
class Net(nn.Module):
    def __init__(s): super().__init__(); s.emb=nn.Embedding(P,D_EMB); s.fc1=nn.Linear(2*D_EMB,D_HID); s.fc2=nn.Linear(D_HID,P)
    def forward(s,a,b): return s.fc2(F.relu(s.fc1(torch.cat([s.emb(a),s.emb(b)],-1))))
g=torch.Generator().manual_seed(args.seed)
a,b=torch.meshgrid(torch.arange(P),torch.arange(P),indexing='ij'); a,b=a.reshape(-1),b.reshape(-1); y=(a+b)%P
perm=torch.randperm(P*P,generator=g); ntr=int(FRAC*P*P); tr,te=perm[:ntr],perm[ntr:]
torch.manual_seed(args.seed); m=Net(); opt=torch.optim.AdamW(m.parameters(),lr=args.lr,betas=BETAS,weight_decay=0.0,eps=args.eps)
ps=list(m.parameters()); n=sum(p.numel() for p in ps)
def nrm(): return torch.sqrt(sum((p.detach()**2).sum() for p in ps)).item()
@torch.no_grad()
def project():
    s=args.norm/nrm()
    for p in ps: p.mul_(s)
project()
print('step,train_acc,test_acc,rms_grad,rms_mhat,rms_sqrt_vhat,rms_step_over_lr,frac_sat,train_loss')
for step in range(args.steps+1):
    lg=m(a[tr],b[tr]); loss=F.cross_entropy(lg,y[tr]); opt.zero_grad(); loss.backward()
    gn=math.sqrt(sum((p.grad.detach()**2).sum().item() for p in ps)/n)
    before=[p.detach().clone() for p in ps]
    opt.step()
    d1=math.sqrt(sum(((p.detach()-q)**2).sum().item() for p,q in zip(ps,before))/n)/args.lr
    project()
    if step%50==0:
        t=step+1; mh=0; vh=0
        for p in ps:
            st=opt.state[p]; mh+=(st['exp_avg']/(1-BETAS[0]**t)).pow(2).sum().item(); vh+=(st['exp_avg_sq']/(1-BETAS[1]**t)).sum().item()
        with torch.no_grad():
            tra=(m(a[tr],b[tr]).argmax(-1)==y[tr]).float().mean().item(); lt=m(a[te],b[te]); tea=(lt.argmax(-1)==y[te]).float().mean().item()
            pc=F.softmax(m(a[tr],b[tr]),-1).gather(1,y[tr][:,None]).squeeze(1); fs=(pc>1-1e-4).float().mean().item()
        print(f'{step},{tra:.4f},{tea:.4f},{gn:.3e},{math.sqrt(mh/n):.3e},{math.sqrt(vh/n):.3e},{d1:.4f},{fs:.3f},{loss.item():.3e}',flush=True)