import json, math, random
from pathlib import Path
import numpy as np
import torch
from torch import nn

SEED = 2993

def seed_all(seed=SEED):
    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
    if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)

class AdaptiveTR:
    def __init__(self, r0, rmin, rmax, kp=.02, km=.01, ema_decay=.9, eps=1e-12):
        self.r=float(r0); self.rmin=float(rmin); self.rmax=float(rmax)
        self.kp=kp; self.km=km; self.decay=ema_decay; self.eps=eps
        self.q_ema=0.; self.steps=0
    def apply(self, params, lr):
        # params have gradients; proposed update is ordinary SGD p=-lr*g.
        sq=0.;
        for p in params:
            if p.grad is not None: sq += float((lr*p.grad).pow(2).sum().item())
        norm=math.sqrt(sq)
        q=min(1., norm/(self.r+self.eps))
        scale=min(1., self.r/(norm+self.eps))
        with torch.no_grad():
            for p in params:
                if p.grad is not None: p.add_(p.grad, alpha=-lr*scale)
        self.q_ema=self.decay*self.q_ema+(1-self.decay)*q
        self.r=float(np.clip(self.r*math.exp(self.kp*q-self.km*(1-q)),self.rmin,self.rmax))
        self.steps += 1
        return norm,q,scale,self.r

def toy_check():
    # Formula-level bidirectional check: saturated proposals expand; unsaturated ones contract.
    r0=.1; tr=AdaptiveTR(r0,.01,10.,kp=.08,km=.04,ema_decay=0.)
    rs=[]; qs=[]
    for proposal in [.5]*8+[.01]*12:
        # emulate a one-dimensional gradient whose proposed norm is proposal
        p=torch.tensor([proposal], requires_grad=True)
        p.grad=torch.tensor([-proposal]) # lr=1 gives proposed norm proposal
        _,q,_,r=tr.apply([p],1.)
        qs.append(q); rs.append(r)
    expansion=rs[7]>r0 and all(rs[i+1]>rs[i] for i in range(7))
    contraction=rs[-1]<rs[8] and all(rs[i+1]<rs[i] for i in range(8,len(rs)-1))
    # Verify smooth quadratic descent inequality in unsaturated regime.
    L=3.; eta=.4; x=torch.tensor([1.2,-.7]); g=L*x
    old=.5*L*float((x*x).sum()); y=x-eta*g
    new=.5*L*float((y*y).sum()); bound=old-(eta-L*eta*eta/2)*float((g*g).sum())
    inequality_ok=new <= bound+1e-6
    return {'radius_start':r0,'radius_after_saturated':rs[7],'radius_end':rs[-1],
            'q_saturated':qs[0],'q_unsaturated':qs[-1],
            'expands_when_saturated':expansion,'contracts_when_unsaturated':contraction,
            'quadratic_descent_inequality_ok':inequality_ok,'quadratic_loss_old':old,'quadratic_loss_new':new,'quadratic_bound':bound}

def make_data(n=2400):
    g=np.random.default_rng(SEED)
    x=g.normal(size=(n,2)).astype('float32')
    y=((x[:,0]*x[:,1]>0).astype('int64'))
    # modest label noise makes the task nontrivial
    flip=g.random(n)<.04; y[flip]=1-y[flip]
    return torch.tensor(x),torch.tensor(y)

class MLP(nn.Module):
    def __init__(self):
        super().__init__(); self.net=nn.Sequential(nn.Linear(2,32),nn.Tanh(),nn.Linear(32,32),nn.Tanh(),nn.Linear(32,2))
    def forward(self,x): return self.net(x)

def run(mode, x, y, steps=500, batch=64):
    seed_all(SEED+17) # identical initialization for both methods
    model=MLP(); lossfn=nn.CrossEntropyLoss()
    # r0 intentionally causes early saturation; fixed control uses same safety cap.
    lr=.35; r0=.055
    tr=AdaptiveTR(r0,.005,.55,kp=.04,km=.02,ema_decay=.9) if mode=='adaptive' else None
    losses=[]; accs=[]; qs=[]; radii=[]; scales=[]
    gen=torch.Generator().manual_seed(SEED+99)
    for t in range(steps):
        ix=torch.randint(0,len(x),(batch,),generator=gen)
        model.zero_grad(set_to_none=True); out=model(x[ix]); loss=lossfn(out,y[ix]); loss.backward()
        if mode=='adaptive': norm,q,scale,r=tr.apply(list(model.parameters()),lr); radii.append(r)
        else:
            sq=sum(float((lr*p.grad).pow(2).sum()) for p in model.parameters() if p.grad is not None)
            norm=math.sqrt(sq); q=min(1.,norm/(r0+1e-12)); scale=min(1.,r0/(norm+1e-12))
            with torch.no_grad():
                for p in model.parameters():
                    if p.grad is not None: p.add_(p.grad,alpha=-lr*scale)
        losses.append(float(loss)); qs.append(q); scales.append(scale)
        if (t+1)%50==0:
            with torch.no_grad(): accs.append(float((model(x).argmax(1)==y).float().mean()))
    with torch.no_grad(): final_loss=float(lossfn(model(x),y)); final_acc=float((model(x).argmax(1)==y).float().mean())
    return {'final_loss':final_loss,'final_accuracy':final_acc,'mean_q':float(np.mean(qs)),
            'clip_fraction':float(np.mean(np.array(qs)>=1-1e-9)),'mean_scale':float(np.mean(scales)),
            'initial_q':float(qs[0]),'last100_q':float(np.mean(qs[-100:])),
            'radius_initial':r0,'radius_final':float(radii[-1]) if radii else r0,
            'loss_at_100':losses[99],'loss_at_500':losses[-1]}

def main():
    seed_all(); toy=toy_check(); x,y=make_data()
    results={'toy_check':toy,'adaptive':run('adaptive',x,y),'fixed_clip':run('fixed',x,y)}
    Path('results.json').write_text(json.dumps(results,indent=2))
    print(json.dumps(results,indent=2))
if __name__=='__main__': main()