"""sgd_clamp.py: copy of ../temp-clamp/temp_clamp.py adding --opt adam|sgd|sgd_norm (sgd: momentum 0.9, no wd; sgd_norm: grad scaled to unit global L2 norm before the momentum-0.9 SGD step). Tag includes the optimizer.
temp_clamp.py was a copy of ../layer-clamp/layer_clamp.py extended with logit_std, frac_sat, grad_norm logging and early stopping.
Per-layer norm clamp on the (a+b) mod 97 MLP (same model/data/optimizer as ../norm-clamp/clamp.py, wd=0).
--clamp total|emb|fc1|fc2|none : which parameter group is rescaled to --norm after every step.
Logs per-layer norms (emb, fc1 incl. bias, fc2 incl. bias) every 50 steps to logs/<tag>.csv and saves logs/<tag>.pt.
Usage: python layer_clamp.py --clamp emb --norm 50 [--seed 0] [--steps 4000] [--lr 1e-3] [--temp T]  (logits are divided by T)
"""
import argparse, csv, os, time, torch
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; BETAS=(0.9,0.98); LOG=50
ap=argparse.ArgumentParser(); ap.add_argument('--clamp',default='total'); ap.add_argument('--norm',type=float,default=0)
ap.add_argument('--seed',type=int,default=0); ap.add_argument('--steps',type=int,default=4000); ap.add_argument('--lr',type=float,default=1e-3); ap.add_argument('--temp',type=float,default=1.0); ap.add_argument('--patience',type=int,default=6); ap.add_argument('--eps',type=float,default=1e-8); ap.add_argument('--wd',type=float,default=0.0); ap.add_argument('--opt',default='adam',choices=['adam','sgd','sgd_norm']); args=ap.parse_args()
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))))/args.temp
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=args.wd,eps=args.eps) if args.opt=='adam' else torch.optim.SGD(m.parameters(),lr=args.lr,momentum=0.9,weight_decay=args.wd)
groups={'total':list(m.parameters()),'emb':[m.emb.weight],'fc1':list(m.fc1.parameters()),'fc2':list(m.fc2.parameters()),'none':[]}
def nrm(ps): return torch.sqrt(sum((p.detach()**2).sum() for p in ps)).item() if ps else 0.0
@torch.no_grad()
def project():
    ps=groups[args.clamp]
    if not ps or args.norm<=0: return
    s=args.norm/nrm(ps)
    for p in ps: p.mul_(s)
@torch.no_grad()
def ev(idx):
    lg=m(a[idx],b[idx]); return F.cross_entropy(lg,y[idx]).item(),(lg.argmax(-1)==y[idx]).float().mean().item(),lg.abs().max().item()
@torch.no_grad()
def train_stats():
    lg=m(a[tr],b[tr]).float(); p=F.softmax(lg,-1).gather(1,y[tr][:,None]).squeeze(1)
    return lg.std().item(),(p>1-1e-4).float().mean().item()
init=[nrm(groups[k]) for k in ('total','emb','fc1','fc2')]
project()
tag=f"{args.opt}_{args.clamp}{args.norm:g}_lr{args.lr:g}_T{args.temp:g}_seed{args.seed}"+("" if args.eps==1e-8 else f"_eps{args.eps:g}")+("" if args.wd==0 else f"_wd{args.wd:g}"); os.makedirs('logs',exist_ok=True)
f=open(f'logs/{tag}.csv','w',newline=''); w=csv.writer(f); w.writerow(['step','train_loss','test_loss','train_acc','test_acc','norm_total','norm_emb','norm_fc1','norm_fc2','max_logit','logit_std','frac_sat','grad_norm'])
t0=time.time(); streak=0
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=torch.sqrt(sum((p.grad.detach()**2).sum() for p in m.parameters())).item()  # raw grad norm before the step
    if step%LOG==0:
        trl,tra,_=ev(tr); tel,tea,ml=ev(te); ls,fs=train_stats()
        w.writerow([step,f'{trl:.6g}',f'{tel:.6g}',f'{tra:.4f}',f'{tea:.4f}']+[f'{nrm(groups[k]):.2f}' for k in ('total','emb','fc1','fc2')]+[f'{ml:.2f}',f'{ls:.4g}',f'{fs:.4f}',f'{gn:.4g}']); f.flush()
        streak=streak+1 if tea>=0.99 else 0
        if streak>=args.patience or step==args.steps: break
    if not torch.isfinite(loss): print('NaN/inf loss at step',step); break
    if args.opt=='sgd_norm':
        with torch.no_grad():
            for p in m.parameters(): p.grad.mul_(1.0/(gn+1e-30))
    opt.step(); project()
torch.save(m.state_dict(),f'logs/{tag}.pt')
print(f"{tag}: init norms total/emb/fc1/fc2={['%.1f'%x for x in init]} final={['%.1f'%nrm(groups[k]) for k in ('total','emb','fc1','fc2')]} train_acc={tra:.3f} test_acc={tea:.3f} max_logit={ml:.1f} last_step={step} wall={time.time()-t0:.0f}s")