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

SEED = 3081
random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
torch.set_num_threads(4)
device = 'cuda' if torch.cuda.is_available() else 'cpu'
try:
    if device == 'cuda': torch.cuda.get_device_name(0)
except Exception:
    device = 'cpu'

# Periodic advection-diffusion: x'(t)=A x(t).
d = 8
L = np.zeros((d,d), np.float32); D = np.zeros((d,d), np.float32)
for i in range(d):
    L[i,i] = -2; L[i,(i-1)%d] = 1; L[i,(i+1)%d] = 1
    D[i,(i+1)%d] = .5; D[i,(i-1)%d] = -.5
A = .7*L + .35*D
At = torch.tensor(A, device=device)
def exact(x,t): return x @ torch.matrix_exp(At.T*float(t))
def rmse(x): return float(torch.sqrt(torch.mean(x*x)).detach().cpu())

def math_check():
    torch.manual_seed(SEED+1); x=torch.randn(64,d,device=device); gen=x@At.T
    gen_errors=[]; growth=[]
    for h in [.2,.1,.05,.025]:
        gen_errors.append(rmse((x+h*gen-x)/h-gen))
        B=torch.randn(d,d,device=device); B=B/torch.linalg.matrix_norm(B)*.002
        ref=x.clone(); bad=x.clone(); errs=[]
        for _ in range(round(1/h)):
            ref=ref+h*(ref@At.T); bad=bad+h*(bad@At.T)+bad@B.T; errs.append(rmse(bad-ref))
        # Measured induced norm gives the discrete Gronwall bound.
        Q=torch.eye(d,device=device)+h*At+B
        Lh=float(torch.linalg.matrix_norm(Q,ord=2).detach().cpu())
        eps=.002*float(torch.linalg.matrix_norm(x,ord=2).detach().cpu())
        bound=eps*sum(Lh**j for j in range(round(1/h)))
        growth.append({'h':h,'final_error':errs[-1],'bound':bound,'Lh':Lh})
    return {'generator_errors':gen_errors,'growth':growth}

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

def rollout(model,x,h,K):
    y=x
    for _ in range(K): y=y+h*model(y)
    return y

def train_model(model, tied=True, steps=500):
    opt=torch.optim.Adam(model.parameters(),lr=3e-3)
    torch.manual_seed(SEED+4)
    x=torch.randn(512,d,device=device)
    for it in range(steps):
        ix=torch.randint(0,x.shape[0],(64,),device=device); z=x[ix]
        # Train all models on the same physical horizon t=1.
        h=.05; K=20; target=exact(z,1.0)
        if tied:
            pred=rollout(model,z,h,K)
            one=(model(z)-((exact(z,.05)-z)/.05)).pow(2).mean()
        else:
            y=z
            for layer in model: y=y+h*layer(y)
            pred=y; one=0.0
        loss=((pred-target)**2).mean()+(.1*one if tied else 0)
        opt.zero_grad(); loss.backward(); opt.step()
    return model

def main():
    check=math_check(); torch.manual_seed(SEED+2)
    test=torch.randn(256,d,device=device)
    tied=train_model(Derivative().to(device),True)
    torch.manual_seed(SEED+3)
    untied=train_model(nn.ModuleList([Derivative() for _ in range(20)]).to(device),False)
    results={}
    # Untied baseline is a conventional depth-20 network trained for t=1.
    with torch.no_grad():
        yt=rollout(tied,test,.05,20)
        yu=test.clone()
        for layer in untied: yu=yu+.05*layer(yu)
        target=exact(test,1)
        results['1']={'tied':rmse(yt-target),'untied':rmse(yu-target)}
        for T in [2,5,10]:
            results[str(T)]={'tied':rmse(rollout(tied,test,.05,round(T/.05))-exact(test,T))}
    out={'device':device,'math':check,'params_tied':sum(p.numel() for p in tied.parameters()),'params_untied':sum(p.numel() for p in untied.parameters()),'rollout':results}
    print(json.dumps(out,indent=2))
    with open('results.json','w') as f: json.dump(out,f,indent=2)
if __name__=='__main__': main()
if __name__=='__main__': main()