"""clamp_gen.py (run 5): generalization of ../temp-clamp/temp_clamp.py.
Fixed-total-norm MLP for (a+b) mod P with selectable modulus, Adam betas, eps, and optimizer variant.
New flags vs temp_clamp.py:
  --p P            modulus (default 97); train fraction stays 0.5 of P*P pairs
  --beta2 B        Adam beta2 (default 0.98); beta1 fixed 0.9
  --opt adamw|atan2   atan2 = Adam-atan2 (Everett et al. 2024): update = a*atan2(m_hat, b*sqrt(v_hat)), no eps
  --a, --b         Adam-atan2 constants (defaults 1.0, 1.0; paper experiments use a=b=1, 1.27=1/arctan(1) is from its Table C4); atan2 tags carry _a{a}_b{b}
  --d_hid H, --d_emb E  hidden width (default 256) and per-token embedding dim (default 64); tags carry _h{H}/_e{E} only when non-default
  --clamp total|none, --norm N, --wd, --lr, --temp, --eps, --f64, --steps, --patience, --seed as before
Logs every 50 steps to logs/<tag>.csv: step, train_loss, test_loss, train_acc, test_acc, norm_total, norm_emb, norm_fc1, norm_fc2,
max_logit(test), logit_std(train), frac_sat(train, p_correct > 1-1e-4), grad_norm (raw, before step), grad_rms (per-coordinate RMS).
Tag: p{P}_{clamp}{N}_lr{lr}_T{T}_b2{beta2}_{opt}[_eps..][_wd..][_f64][_frac..][_h..][_e..]_seed{s}
"""
import argparse, csv, os, time, math, torch
torch.set_num_threads(1)
import torch.nn as nn, torch.nn.functional as F
ap=argparse.ArgumentParser()
ap.add_argument('--p',type=int,default=97); 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=8000); 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('--f64',action='store_true'); ap.add_argument('--beta2',type=float,default=0.98)
ap.add_argument('--opt',default='adamw'); ap.add_argument('--a',type=float,default=1.0); ap.add_argument('--b',type=float,default=1.0)
ap.add_argument('--outdir',default='logs'); ap.add_argument('--frac',type=float,default=0.5)
ap.add_argument('--d_hid',type=int,default=256); ap.add_argument('--d_emb',type=int,default=64)
args=ap.parse_args()
if args.f64: torch.set_default_dtype(torch.float64)
P=args.p; D_EMB,D_HID=args.d_emb,args.d_hid; LOG=50; 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))))/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(args.frac*P*P); tr,te=perm[:ntr],perm[ntr:]
torch.manual_seed(args.seed); m=Net()

class AdamAtan2(torch.optim.Optimizer):
    """Adam-atan2 (Everett et al. 2024, arXiv 2407.05872): p -= lr * a * atan2(m_hat, b*sqrt(v_hat)); decoupled wd as in AdamW."""
    def __init__(s,params,lr,betas,weight_decay,a,b): super().__init__(params,dict(lr=lr,betas=betas,weight_decay=weight_decay,a=a,b=b))
    @torch.no_grad()
    def step(s):
        for grp in s.param_groups:
            b1,b2=grp['betas']
            for p in grp['params']:
                if p.grad is None: continue
                st=s.state[p]
                if not st: st['t']=0; st['m']=torch.zeros_like(p); st['v']=torch.zeros_like(p)
                st['t']+=1; t=st['t']; gr=p.grad
                st['m'].mul_(b1).add_(gr,alpha=1-b1); st['v'].mul_(b2).addcmul_(gr,gr,value=1-b2)
                mh=st['m']/(1-b1**t); vh=st['v']/(1-b2**t)
                if grp['weight_decay']: p.mul_(1-grp['lr']*grp['weight_decay'])
                p.add_(grp['a']*torch.atan2(mh,grp['b']*vh.sqrt()),alpha=-grp['lr'])

if args.opt=='adamw': opt=torch.optim.AdamW(m.parameters(),lr=args.lr,betas=BETAS,weight_decay=args.wd,eps=args.eps)
elif args.opt=='atan2': opt=AdamAtan2(m.parameters(),lr=args.lr,betas=BETAS,weight_decay=args.wd,a=args.a,b=args.b)
else: raise SystemExit('unknown --opt')
groups={'total':list(m.parameters()),'emb':[m.emb.weight],'fc1':list(m.fc1.parameters()),'fc2':list(m.fc2.parameters()),'none':[]}
nparam=sum(p.numel() for p in m.parameters())
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"p{P}_{args.clamp}{args.norm:g}_lr{args.lr:g}_T{args.temp:g}_b2{args.beta2:g}_{args.opt}"
     +("" if (args.eps==1e-8 or args.opt!='adamw') else f"_eps{args.eps:g}")+("" if args.opt!='atan2' else f"_a{args.a:g}_b{args.b:g}")+("" if args.wd==0 else f"_wd{args.wd:g}")
     +("_f64" if args.f64 else "")+("" if args.frac==0.5 else f"_frac{args.frac:g}")+("" if D_HID==256 else f"_h{D_HID}")+("" if D_EMB==64 else f"_e{D_EMB}")+f"_seed{args.seed}")
os.makedirs(args.outdir,exist_ok=True)
f=open(f'{args.outdir}/{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','grad_rms'])
t0=time.time(); streak=0; first99=None
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()
    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'{gn/math.sqrt(nparam):.4g}']); f.flush()
        if tea>=0.99 and first99 is None: first99=step
        streak=streak+1 if tea>=0.99 else 0
        if streak>=args.patience or step==args.steps: break
    opt.step(); project()
torch.save(m.state_dict(),f'{args.outdir}/{tag}.pt')
print(f"{tag}: nparam={nparam} init_norm={init[0]:.1f} final_norm={nrm(groups['total']):.1f} train_acc={tra:.3f} test_acc={tea:.3f} max_logit={ml:.1f} step_test99={first99} last_step={step} wall={time.time()-t0:.0f}s",flush=True)