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

SEED=3010
np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
rng=np.random.default_rng(SEED)
try:
    device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    if device.type=='cuda': torch.cuda.set_device(0)
except Exception: device=torch.device('cpu')

# Bistable scalar dynamics: x=+-1 attractors, x=0 repeller.
def dyn(x): return x + 0.18*x*(1-x*x)
N=rng.normal(size=(1,16)).astype('float32')
def obs(x):
    x=np.asarray(x).reshape(-1,1)
    q=np.concatenate([x,x*x,np.sin(x),np.cos(x)],1).astype('float32')
    y=np.zeros((len(x),20),dtype='float32'); y[:,:4]=q; y[:,4:]=.15*q[:,[0]]@N
    return y

def samples(n):
    x=rng.uniform(-1.45,1.45,n); xn=dyn(x)
    return torch.tensor(obs(x)),torch.tensor(obs(xn)),x,xn

class WM(nn.Module):
    def __init__(self):
        super().__init__()
        self.e=nn.Sequential(nn.Linear(20,32),nn.Tanh(),nn.Linear(32,1))
        self.d=nn.Sequential(nn.Linear(1,32),nn.Tanh(),nn.Linear(32,20))
        self.g=nn.Sequential(nn.Linear(1,16),nn.Tanh(),nn.Linear(16,1))
    def forward(self,x):
        z=self.e(x); return z,self.d(z),self.g(z)

def train(reg):
    tr=samples(2400); va=samples(1000)
    model=WM().to(device); opt=torch.optim.Adam(model.parameters(),lr=2e-3)
    X,Xn=tr[0].to(device),tr[1].to(device)
    for _ in range(700):
        z,rec,p=model(X); zn=model.e(Xn)
        loss=(rec-X).pow(2).mean()+2*(p-zn).pow(2).mean()
        if reg: loss += 2*(p-zn.detach()).pow(2).mean()
        opt.zero_grad(); loss.backward(); opt.step()
    with torch.no_grad():
        V,Vn=va[0].to(device),va[1].to(device)
        z,rec,p=model(V); zn=model.e(Vn)
        residual=(p-zn).abs().cpu().numpy().ravel()
        z0=model.e(V).cpu().numpy().ravel(); zp=model.g(model.e(V)).cpu().numpy().ravel()
        recon=(rec-V).pow(2).mean().sqrt().item()
    return model, residual,z0,zp,recon

def graph_cert(z,zp,eps,nb=28):
    lo,hi=float(np.percentile(z,1)),float(np.percentile(z,99))
    lo-=.15; hi+=.15; w=(hi-lo)/nb
    edges=[set() for _ in range(nb)]
    for i in range(nb):
        a,b=lo+i*w,lo+(i+1)*w; pts=np.linspace(a,b,9)
        # empirical padded image; cells intersecting image interval receive edges
        y=np.interp(pts,z,zp) if np.all(np.diff(z)>0) else None
        # use model-free nearest validation image samples in each source cell
        sel=(z>=a)&(z<=b)
        if not np.any(sel): continue
        ymin,ymax=zp[sel].min()-eps,zp[sel].max()+eps
        js=np.where((lo+(np.arange(nb)+1)*w>=ymin)&(lo+np.arange(nb)*w<=ymax))[0]
        edges[i].update(js.tolist())
    # Tarjan SCC, then terminal components
    idx=0; stack=[]; on=set(); ids=[-1]*nb; low=[0]*nb; comps=[]
    def go(v):
        nonlocal idx
        ids[v]=low[v]=idx; idx+=1; stack.append(v); on.add(v)
        for q in edges[v]:
            if ids[q]<0: go(q); low[v]=min(low[v],low[q])
            elif q in on: low[v]=min(low[v],ids[q])
        if low[v]==ids[v]:
            c=[]
            while True:
                q=stack.pop(); on.remove(q); c.append(q)
                if q==v: break
            comps.append(c)
    for v in range(nb):
        if ids[v]<0: go(v)
    terminal=[]
    for c in comps:
        if not any(q not in c for v in c for q in edges[v]): terminal.append(c)
    # margin to cell boundaries, estimated from padded source images.
    margin=float('inf')
    for i in range(nb):
        if edges[i]:
            a,b=lo+i*w,lo+(i+1)*w
            vals=zp[(z>=a)&(z<=b)]
            if len(vals):
                margin=min(margin,float(np.min(np.minimum(np.abs(vals-lo),np.abs(vals-hi)))))
    return len(terminal), terminal, w, margin

def math_check():
    # For scalar perturbations, a padded interval contains every true image iff eps>=residual.
    g=np.array([-.8,-.2,.3,.9]); true=g+np.array([.01,-.04,.02,-.03]); r=np.abs(true-g)
    return {'max_residual':float(r.max()),'contain_at_eps_max':bool(np.all(np.abs(true-g)<=r.max()+1e-12)), 'contain_below_max':bool(np.all(np.abs(true-g)<=r.max()*.5))}

def main():
    result={'device':str(device),'math_check':math_check()}
    for name,reg in [('baseline',False),('semiconjugacy',True)]:
        model,r,z,zp,recon=train(reg)
        eps=float(np.max(r)); q=float(np.quantile(r,.95))
        n,comps,w,margin=graph_cert(z,zp,eps)
        result[name]={'max_residual':eps,'q95_residual':q,'reconstruction_rmse':recon,'cell_width':w,'estimated_margin':margin,'terminal_sccs':n,'certificate':bool(eps<w)}
    print(json.dumps(result,indent=2))
    with open('results.json','w') as f: json.dump(result,f,indent=2)
if __name__=='__main__': main()