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

SEED=2805
random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
torch.set_num_threads(min(8, torch.get_num_threads()))
try:
    device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    if device.type=='cuda': torch.zeros(1, device=device)
except Exception:
    device=torch.device('cpu')

def math_checks():
    eps=1e-6
    # Prediction 1: away from zero, ds/da = eps/(a+eps)^2.
    amps=np.array([1e-4,1e-3,1e-2,1e-1,1.0])
    delta=1e-8
    s=lambda x: abs(x)/(abs(x)+eps)
    observed=[(s(a+delta)-s(a))/delta for a in amps]
    predicted=[eps/(a+eps)**2 for a in amps]
    rel=float(np.max(np.abs(np.array(observed)-predicted)/(np.array(predicted)+1e-30)))
    # Prediction 2: for K=rI, rollout norm is exactly r^h times initial norm.
    rs=np.array([0.7,0.9,1.0,1.1,1.3]); h=8; v=np.ones(4)/2
    rollout_obs=[np.linalg.norm((r**h)*v) for r in rs]
    rollout_pred=[(r**h)*np.linalg.norm(v) for r in rs]
    rollout_err=float(np.max(np.abs(np.array(rollout_obs)-rollout_pred)))
    # Prediction 3: for |delta| << a, two-sided support change is
    # 2 eps |delta|/(a+eps)^2. Use a delta that never crosses zero.
    delta2=1e-8
    support_obs=[]; support_pred=[]
    for a in amps:
        obs=abs(s(a+delta2)-s(a))+abs(s(a-delta2)-s(a))
        pr=2*eps/(a+eps)**2*delta2
        support_obs.append(obs); support_pred.append(pr)
    support_rel=float(np.max(np.abs(np.array(support_obs)-support_pred)/(np.array(support_pred)+1e-30)))
    return {'support_derivative':{'amplitudes':amps.tolist(),'observed':observed,'predicted':predicted,'max_relative_error':rel},
            'rollout_scaling':{'r':rs.tolist(),'h':h,'observed_norm':rollout_obs,'predicted_norm':rollout_pred,'max_abs_error':rollout_err},
            'support_perturbation':{'amplitudes':amps.tolist(),'delta':delta2,'observed':support_obs,'predicted':support_pred,'max_relative_error':support_rel}}

def make_data(ntraj=180,T=70):
    allx=[]; labels=[]
    for basin in [-1,1]:
        for _ in range(ntraj//2):
            q=basin+np.random.normal(0,.12); v=np.random.normal(0,.12); traj=[]
            for t in range(T):
                traj.append([q,v]); acc=q-q**3-.32*v; q=q+.08*v; v=v+.08*acc
            allx.append(np.asarray(traj,np.float32)); labels.append(basin)
    return np.asarray(allx),np.asarray(labels)

class Koopman(nn.Module):
    def __init__(self,dz=16):
        super().__init__(); self.enc=nn.Sequential(nn.Linear(2,32),nn.Tanh(),nn.Linear(32,dz)); self.dec=nn.Sequential(nn.Linear(dz,32),nn.Tanh(),nn.Linear(32,2)); self.K=nn.Parameter(torch.eye(dz)+.01*torch.randn(dz,dz))
    def forward(self,x): return self.enc(x)

def train_model(x,sparse=True,eta=.002,steps=700,H=5,dz=16):
    torch.manual_seed(SEED+(1 if sparse else 2)+int(eta*100000)); model=Koopman(dz).to(device); opt=torch.optim.Adam(model.parameters(),lr=2e-3); xt=torch.tensor(x,device=device); n,T,_=x.shape; losses=[]
    for step in range(steps):
        ids=torch.randint(0,n,(64,),device=device); ts=torch.randint(0,T-H,(64,),device=device); win=torch.stack([xt[ids,j] for j in range(H+1)],1); z0=model(win[:,0]); pred=z0; rec=((model.dec(z0)-win[:,0])**2).mean(); lat=roll=sup=0.; prev_z=z0
        for j in range(1,H+1):
            pred=pred@model.K.T; zj=model(win[:,j]); lat+=((pred-zj)**2).mean(); roll+=((model.dec(pred)-win[:,j])**2).mean(); ss=lambda z:z.abs()/(z.abs()+1e-6); sup+=(ss(zj)-ss(prev_z)).abs().mean(); prev_z=zj
        loss=rec+.5*lat/H+roll/H+(.0005*z0.abs().mean() if sparse else 0)+(eta*sup if sparse else 0); opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(),10); opt.step()
        if step%100==0: losses.append(float(loss.detach().cpu()))
    return model,losses

@torch.no_grad()
def evaluate(model,x,labels,H=20,tau=None):
    xt=torch.tensor(x,device=device); z=model(xt[:,0]); pred=z; outs=[]
    for h in range(1,H+1): pred=pred@model.K.T; outs.append(model.dec(pred))
    mse=float(((torch.stack(outs,1)-xt[:,1:H+1])**2).mean().cpu()); zz=model(xt.reshape(-1,2)).reshape(len(x),len(x[0]),-1)
    if tau is None: tau=float(torch.quantile(zz.abs().flatten(),.75).cpu())
    masks=(zz[:,0].abs()>tau).cpu().numpy().astype(np.int8); keys=[tuple(a.tolist()) for a in masks]; uniq={k:i for i,k in enumerate(sorted(set(keys)))}; pg=np.array([uniq[k]%2 for k in keys]); y=(labels>0).astype(int); acc=max(np.mean(pg==y),np.mean(1-pg==y)); ss=zz.abs()/(zz.abs()+1e-6)
    return {'mse20':mse,'tau':tau,'active_fraction':float(masks.mean()),'mask_basin_acc':float(acc),'support_step_change':float((ss[:,1:]-ss[:,:-1]).abs().mean().cpu())}

def main():
    t0=time.time(); checks=math_checks(); train_x,train_y=make_data(180,70); test_x,test_y=make_data(60,70); results={}
    base,_=train_model(train_x,False,eta=0,steps=700); results['dense']=evaluate(base,test_x,test_y)
    sparse,_=train_model(train_x,True,eta=.002,steps=700); results['sparse']=evaluate(sparse,test_x,test_y)
    ablation={}
    for eta in [0.,.001,.005]:
        m,_=train_model(train_x,True,eta=eta,steps=450); ablation[str(eta)]=evaluate(m,test_x,test_y)
    report={'seed':SEED,'device':str(device),'math_checks':checks,'baseline_vs_idea':results,'eta_ablation':ablation,'runtime_sec':time.time()-t0,'predictions':['support derivative scales as epsilon/(a+epsilon)^2','linear rollout norm scales exactly as spectral radius^h','small support perturbations scale as 2 epsilon |delta|/(a+epsilon)^2']}
    Path('results.json').write_text(json.dumps(report,indent=2)); print(json.dumps(report,indent=2))
if __name__=='__main__': main()