import json, random
import numpy as np


def toy_checks():
    out = {}; H = 8
    av = np.array([.80,.95,.99,1.,1.01,1.05,1.20]); obs=[]
    for a in av:
        z=g=1.
        for _ in range(H): z*=abs(a); g=max(g,z)
        obs.append(g)
    out['scalar_boundary']={'H':H,'predicted_boundary_abs_a':1.,'sweep_a':av.tolist(),
      'observed_gain':obs,'predicted_gain':[max(1.,abs(float(a))**H) for a in av]}
    Ks=np.array([.25,.75,1.,2.,4.]); trans=[]; final=[]
    for K in Ks:
        A=np.array([[0.,K],[0.,0.]]); d=np.array([0.,1.]); tr=[d.copy()]
        for _ in range(2): d=A@d; tr.append(d.copy())
        trans.append(max(np.linalg.norm(q) for q in tr)); final.append(np.linalg.norm(tr[-1]))
    out['transient_blind_spot']={'K':Ks.tolist(),'predicted_DH_over_eps':[max(1.,float(k)) for k in Ks],
      'observed_DH_over_eps':trans,'final_step_gain':final}
    rho,H2=.8,5; K2=np.array([0.,.5,1.,2.,4.]); gains=[]; pred=[]
    for K in K2:
        A=np.array([[rho,K],[0.,rho]]); d=np.array([0.,1.]); peak=1.
        for k in range(1,H2+1):
            d=A@d; peak=max(peak,np.linalg.norm(d))
        gains.append(peak); pred.append(max([1.]+[float(np.linalg.norm([k*K*rho**(k-1),rho**k])) for k in range(1,H2+1)]))
    out['nonnormal_scaling']={'rho':rho,'H':H2,'K':K2.tolist(),'observed_gain':gains,
      'predicted_gain':pred,'slope_observed_large_K':float((gains[-1]-gains[1])/(K2[-1]-K2[1]))}
    return out


def run_gru(device):
    import torch
    import torch.nn as nn
    torch.manual_seed(7); np.random.seed(7); random.seed(7)
    T,N,B,inp,hid=18,512,64,3,16
    x=torch.randn(N,T,inp,device=device); y=torch.roll(x,-1,1); y[:,-1]=0
    class Model(nn.Module):
        def __init__(self):
            super().__init__(); self.rnn=nn.GRU(inp,hid,batch_first=True); self.head=nn.Linear(hid,inp)
        def forward(self,z,h=None):
            q,h=self.rnn(z,h); return self.head(q),h,q
    def train(lam):
        torch.manual_seed(11); m=Model().to(device); opt=torch.optim.Adam(m.parameters(),lr=.01)
        for step in range(100):
            ix=torch.randint(0,N,(B,),device=device); z=x[ix]; target=y[ix]; p,_,_=m(z)
            task=(p-target).pow(2).mean(); eps=torch.randn(B,hid,device=device)*.02
            h0=torch.zeros(1,B,hid,device=device); h1=h0.clone(); h1[0]+=eps; u=z[:,:6]
            _,_,a=m(u,h0); _,_,b=m(u,h1)
            roll=torch.sqrt((a-b).pow(2).sum(-1)+1e-12).amax(1)/(eps.norm(dim=1)+1e-8)
            loss=task+lam*roll.mean(); opt.zero_grad(); loss.backward(); opt.step()
        with torch.no_grad():
            ix=torch.arange(128,device=device); p,_,_=m(x[ix]); mse=(p-y[ix]).pow(2).mean().item()
            eps=torch.randn(128,hid,device=device)*.02; h0=torch.zeros(1,128,hid,device=device); h1=h0.clone(); h1[0]+=eps
            _,_,a=m(x[ix,:8],h0); _,_,b=m(x[ix,:8],h1)
            gain=(torch.sqrt((a-b).pow(2).sum(-1)+1e-12).amax(1)/eps.norm(dim=1)).median().item()
        return {'loss':mse,'median_G8':gain}
    return {'device':device,'baseline':train(0.),'rollout_lambda_0.03':train(.03)}


def gru_experiment():
    try:
        import torch
        requested='cuda' if torch.cuda.is_available() else 'cpu'
        try: return run_gru(requested)
        except Exception as first:
            if requested=='cuda':
                return {'device':'cpu_fallback','cuda_error':repr(first),'result':run_gru('cpu')}
            return {'error':repr(first)}
    except Exception as e: return {'error':repr(e)}

if __name__=='__main__':
    result={'toy':toy_checks(),'gru':gru_experiment()}
    with open('results.json','w') as f: json.dump(result,f,indent=2)
    print(json.dumps(result,indent=2))