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

SEED = 3142

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)

def device_or_cpu():
    if torch.cuda.is_available():
        try:
            x = torch.zeros(1, device='cuda'); del x
            return torch.device('cuda')
        except Exception as e:
            print('CUDA fallback:', repr(e))
    return torch.device('cpu')

class PlainMLP(nn.Module):
    def __init__(self, d, width=64, depth=3):
        super().__init__()
        layers=[nn.Linear(d,width), nn.Tanh()]
        for _ in range(depth-1): layers += [nn.Linear(width,width), nn.Tanh()]
        layers += [nn.Linear(width,1)]
        self.net=nn.Sequential(*layers)
    def forward(self,x): return self.net(x).squeeze(-1)

class AccumulatorResNet(nn.Module):
    """Persistent context h and additive scalar accumulator s.
    alpha=0 transmits h exactly; each branch adds one correction to s.
    """
    def __init__(self,d,q=32,r=1,K=4,width=48):
        super().__init__(); self.K=K; self.q=q; self.r=r
        self.embed=nn.Sequential(nn.Linear(d,q),nn.Tanh())
        self.s0=nn.Parameter(torch.zeros(r))
        self.branches=nn.ModuleList()
        for _ in range(K):
            self.branches.append(nn.Sequential(nn.Linear(q+r,width),nn.Tanh(),
                                               nn.Linear(width,q+r),nn.Tanh()))
        self.head=nn.Sequential(nn.Linear(q+r,width),nn.Tanh(),nn.Linear(width,1))
    def forward(self,x):
        h=self.embed(x); s=self.s0.expand(x.shape[0],-1)
        for branch in self.branches:
            dh,ds=branch(torch.cat([h,s],1)).split([self.q,self.r],1)
            # paper-motivated gates: context is carried, accumulator adds correction
            h = h + 0.0*dh
            s = s + ds
        return self.head(torch.cat([h,s],1)).squeeze(-1)

class StandardResNet(nn.Module):
    """Control: same state and branches, but updates all coordinates residually."""
    def __init__(self,d,q=32,r=1,K=4,width=48):
        super().__init__(); self.K=K; self.q=q; self.r=r
        self.embed=nn.Sequential(nn.Linear(d,q),nn.Tanh()); self.s0=nn.Parameter(torch.zeros(r))
        self.branches=nn.ModuleList([nn.Sequential(nn.Linear(q+r,width),nn.Tanh(),
            nn.Linear(width,q+r),nn.Tanh()) for _ in range(K)])
        self.head=nn.Sequential(nn.Linear(q+r,width),nn.Tanh(),nn.Linear(width,1))
    def forward(self,x):
        z=torch.cat([self.embed(x),self.s0.expand(x.shape[0],-1)],1)
        for b in self.branches: z=z+b(z)
        return self.head(z).squeeze(-1)

def target(x):
    d=x.shape[1]
    # A high-dimensional, bounded analogue of a heat/Picard target: additive
    # first-order term plus a smooth nonlinear correction.
    a=torch.arange(1,d+1,device=x.device,dtype=x.dtype)
    return (torch.sin(x)*((a%7+1)/4)).sum(1)/math.sqrt(d) + 0.25*torch.tanh(x.sum(1)/math.sqrt(d))

def count(m): return sum(p.numel() for p in m.parameters() if p.requires_grad)

def math_check():
    rng=np.random.default_rng(7); n=30; q=5; r=2
    h=rng.normal(size=(n,q)); s=rng.normal(size=(n,r)); original=s.copy(); ds=[]
    for k in range(6):
        delta=rng.normal(size=(n,r))/(k+1); ds.append(delta); s=s+delta
    err=float(np.max(np.abs(s-(original+sum(ds)))))
    # Explicit composition of two additive blocks equals one concatenated rollout.
    x=rng.normal(size=(n,3)); A=rng.normal(size=(3,2)); B=rng.normal(size=(2,2)); C=rng.normal(size=(2,2))
    y=x@A; y1=y+y@B; y2=y1+y1@C
    composed=(x@A)+((x@A)@B)+((x@A+(x@A)@B)@C)
    return {'telescoping_max_error':err,'composition_max_error':float(np.max(np.abs(y2-composed))),
            'bounded_tanh_max_abs':float(np.max(np.abs(np.tanh(rng.normal(size=10000))))) }

def make_data(d, n, dev, seed):
    g=torch.Generator(device='cpu').manual_seed(seed)
    x=torch.randn(n,d,generator=g)
    return x.to(dev), target(x.to(dev))

def train(model,x,y,xt,yt,steps=450,batch=128):
    opt=torch.optim.Adam(model.parameters(),lr=2e-3); N=len(x); hist=[]
    model.train(); t0=time.time()
    for step in range(steps):
        ix=torch.randint(N,(min(batch,N),),device=x.device)
        loss=((model(x[ix])-y[ix])**2).mean(); opt.zero_grad(); loss.backward(); opt.step()
        if step in (0,99,249,449):
            model.eval()
            with torch.no_grad(): hist.append(float(((model(xt)-yt)**2).mean().cpu()))
            model.train()
    model.eval()
    with torch.no_grad(): mse=float(((model(xt)-yt)**2).mean().cpu())
    return mse,hist,time.time()-t0

def main():
    seed_all(); dev=device_or_cpu(); print('device',dev)
    result={'device':str(dev),'math_check':math_check(),'runs':[]}
    for d in (50,100,200):
        x,y=make_data(d,1800,dev,100+d); xt,yt=make_data(d,600,dev,900+d)
        for name, cls in [('plain_mlp',PlainMLP),('standard_resnet',StandardResNet),('accumulator',AccumulatorResNet)]:
            seed_all(SEED+d+len(name)); m=cls(d).to(dev)
            try: mse,hist,secs=train(m,x,y,xt,yt)
            except Exception as e:
                if dev.type=='cuda':
                    print('CUDA error, rerun CPU:',repr(e)); dev=torch.device('cpu'); x,y,xt,yt=[z.cpu() for z in (x,y,xt,yt)]; m=cls(d).to(dev); mse,hist,secs=train(m,x,y,xt,yt)
                else: raise
            result['runs'].append({'d':d,'model':name,'params':count(m),'test_mse':mse,'checkpoints':hist,'seconds':secs})
            print(d,name,count(m),mse,hist)
    Path('results.json').write_text(json.dumps(result,indent=2))
    print('wrote results.json')
if __name__=='__main__': main()