import json, random, math
import numpy as np
import torch
import torch.nn as nn
from torch.nn.utils.parametrizations import spectral_norm

SEED=2975
random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
device='cuda' if torch.cuda.is_available() else 'cpu'
torch.set_default_dtype(torch.float64)

class Block(nn.Module):
    def __init__(self, d=2, hidden=16, eta=.15, alpha=.20, idea=False):
        super().__init__(); self.d=d; self.eta=eta; self.alpha=alpha; self.idea=idea
        self.l1=nn.Linear(2*d,hidden); self.l2=nn.Linear(hidden,d)
        if idea:
            # tanh and the concatenation projection are 1-Lipschitz. Spectral
            # normalization makes each learned linear map 1-Lipschitz.
            self.l1=spectral_norm(self.l1); self.l2=spectral_norm(self.l2)
            # Reserve a Lipschitz budget for the residual correction.
            self.skip=math.sqrt(1-alpha)-eta
        else: self.skip=1.0
    def forward(self,z,u):
        h=torch.tanh(self.l1(torch.cat([z,u],-1)))
        return self.skip*z+self.eta*self.l2(h)

def jac(model,z,u):
    z=z.detach().requires_grad_(True); y=model(z[None],u[None])[0]
    return torch.stack([torch.autograd.grad(y[i],z,retain_graph=True)[0] for i in range(y.numel())])
def eig(a): return torch.linalg.eigvalsh((a+a.T)/2)[-1]
def report(model, alpha=.2):
    g=torch.Generator(device=device).manual_seed(SEED+11)
    zs=torch.rand(64,2,generator=g,device=device)*2-1; us=torch.rand(64,2,generator=g,device=device)*2-1
    ev=[float(eig(jac(model,z,u).T@jac(model,z,u)-(1-alpha)*torch.eye(2,device=device)).detach().cpu()) for z,u in zip(zs,us)]
    z1=torch.tensor([.65,-.45],device=device); z2=z1+torch.tensor([1e-3,-1.2e-3],device=device); u=torch.tensor([.2,-.1],device=device)
    ratios=[]
    for _ in range(20):
        old=torch.linalg.vector_norm(z2-z1); z1=model(z1[None],u[None])[0]; z2=model(z2[None],u[None])[0]
        ratios.append(float((torch.linalg.vector_norm(z2-z1)/old).detach().cpu()))
    return {'max_violation':max(ev),'mean_violation':float(np.mean(ev)), 'max_ratio':max(ratios),'mean_ratio':float(np.mean(ratios)), 'distance_ratio_20':float(np.prod(ratios)), 'ratios':ratios}

def train(model, constrained, epochs=250):
    g=torch.Generator(device=device).manual_seed(SEED)
    z=torch.rand(96,2,generator=g,device=device)*2-1; u=torch.rand(96,2,generator=g,device=device)*2-1
    target=1.08*z+.25*u+.04*torch.sin(2*z)
    opt=torch.optim.Adam(model.parameters(),lr=3e-3)
    for _ in range(epochs):
        opt.zero_grad(); pred=model(z,u); task=((pred-target)**2).mean()
        # Exact Jacobian penalty, retained for the idea even though its
        # spectral-normalized budget already gives a structural certificate.
        qs=[]
        for k in range(0,96,8):
            M=jac(model,z[k],u[k]); qs.append(torch.nn.functional.softplus(eig(M.T@M-(1-model.alpha)*torch.eye(2,device=device))+.02)**2)
        penalty=torch.stack(qs).mean(); (task+(2*penalty if constrained else 0)).backward(); opt.step()
    return float(task.detach().cpu()),float(penalty.detach().cpu())

def main():
    global device
    # Algebra sanity check: the discrete LMI is exactly equivalent here to
    # the largest eigenvalue of M^T M-(1-alpha)I being nonpositive.
    good=.8*torch.eye(2); bad=1.1*torch.eye(2); I=torch.eye(2)
    out={'seed':SEED,'device':device,'math_check':{
      'alpha':.3,'good_largest_eigenvalue':float(eig(good.T@good-(1-.3)*I)),
      'bad_largest_eigenvalue':float(eig(bad.T@bad-(1-.3)*I)),
      'good_distance_after_12':.8**12,'bad_distance_after_12':1.1**12}}
    for name,flag in [('baseline',False),('idea',True)]:
        try:
            m=Block(idea=flag).to(device); loss,pen=train(m,flag); out[name]={'task_mse':loss,'penalty':pen,**report(m)}
        except Exception:
            if device!='cuda': raise
            device='cpu'; m=Block(idea=flag); loss,pen=train(m,flag); out[name]={'task_mse':loss,'penalty':pen,**report(m)}
    with open('safe_results.json','w') as f: json.dump(out,f,indent=2)
    print(json.dumps(out,indent=2))
if __name__=='__main__': main()