import os, json, math, random
import numpy as np
import torch
import torch.nn as nn

SEED=17
random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
torch.set_num_threads(4)
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 sym_eigs(M): return torch.linalg.eigvalsh((M+M.transpose(-1,-2))/2)

# Core numerical verification: direct finite-step LMI and energy ratio.
def exact_check():
    h=.08; alpha=.35
    # A is stable in a non-Euclidean metric; P changes with theta.
    P0=torch.tensor([[2.0,.35],[.35,1.0]])
    P1=P0.clone()
    # Construct A from a contracting similarity transform, then test the actual discrete condition.
    Q=torch.tensor([[-1.4, .25],[-.10,-.9]])
    L=torch.linalg.cholesky(P0)
    A=torch.linalg.solve(L.T, Q @ L.T) # gives A^T P + P A = L(Q^T+Q)L^T
    J=torch.eye(2)+h*A
    S=J.T@P1@J-math.exp(-2*alpha*h)*P0
    # Generalized normalized matrix via symmetric inverse square root
    ew,U=torch.linalg.eigh(P0); Pinv=(U*ew.rsqrt())@U.T
    lam=torch.linalg.eigvalsh(Pinv@S@Pinv).max().item()
    # simulate perturbations and compare exact metric energies
    d=torch.tensor([1.0,-.7]); v0=(d@P0@d).item(); d1=J@d; v1=(d1@P1@d1).item()
    ratio=v1/v0; bound=math.exp(-2*alpha*h)
    # A deliberately noncontracting control in same P.
    Jbad=torch.eye(2)+h*(A+1.8*torch.eye(2)); Sbad=Jbad.T@P1@Jbad-bound*P0
    lambad=torch.linalg.eigvalsh(Pinv@Sbad@Pinv).max().item()
    return {'device':str(device),'h':h,'alpha':alpha,'finite_lmi_lambda':lam,
            'energy_ratio':ratio,'bound':bound,'bad_lambda':lambad,
            'lmi_satisfied':bool(lam<=1e-7),'energy_satisfied':bool(ratio<=bound)}

class Dyn(nn.Module):
    def __init__(self, hidden=32):
        super().__init__()
        self.net=nn.Sequential(nn.Linear(3,hidden),nn.Tanh(),nn.Linear(hidden,hidden),nn.Tanh(),nn.Linear(hidden,2))
        self.pnet=nn.Sequential(nn.Linear(1,16),nn.Tanh(),nn.Linear(16,3))
    def forward(self,z,theta): return z + .12*self.net(torch.cat([z,theta],-1))
    def metric(self,theta):
        x=self.pnet(theta); L=torch.zeros(x.shape[0],2,2,device=x.device)
        L[:,0,0]=torch.nn.functional.softplus(x[:,0])+0.15
        L[:,1,0]=x[:,1]; L[:,1,1]=torch.nn.functional.softplus(x[:,2])+0.15
        return L@L.transpose(-1,-2)+.05*torch.eye(2,device=x.device)

def target(z,theta):
    # Two changing operating regimes, both stable but with different shear/rotation.
    th=theta[:,0]
    a=0.73+0.10*torch.sin(th); b=.18*torch.cos(th)
    M=torch.stack([torch.stack([a,b],-1),torch.stack([-b,a-.08*torch.sin(th)],-1)],-2)
    return torch.bmm(M,z.unsqueeze(-1)).squeeze(-1)

def jacobian_batch(model,z,theta):
    # Explicit small-state Jacobians, retaining autograd for the penalty.
    outs=[]
    for i in range(z.shape[0]):
        zi=z[i:i+1].detach().requires_grad_(True); ti=theta[i:i+1]
        yi=model(zi,ti).squeeze(0)
        rows=[]
        for q in range(2): rows.append(torch.autograd.grad(yi[q],zi,create_graph=True,retain_graph=True)[0].squeeze(0))
        outs.append(torch.stack(rows))
    return torch.stack(outs)

def lyap_penalty(model,z,theta,alpha=.22,h=.12):
    J=jacobian_batch(model,z,theta); P=model.metric(theta)
    # P_{k+1}: parameter changes between adjacent operating conditions (same batch shuffled partner).
    theta1=torch.roll(theta,1,0); P1=model.metric(theta1)
    E=torch.eye(2,device=z.device).expand(z.shape[0],-1,-1)
    R=J.transpose(1,2)@P1@J-math.exp(-2*alpha*h)*P
    # generalized max eigenvalue, normalized by P^{-1/2}; squared positive part.
    ev,U=torch.linalg.eigh(P); Pinv=(U*ev.rsqrt().unsqueeze(-2))@U.transpose(1,2)
    lam=torch.linalg.eigvalsh(Pinv@R@Pinv)[:,-1]
    return torch.relu(lam).square().mean(), lam.detach()

def run(penalized, steps=120):
    torch.manual_seed(SEED+int(penalized)); model=Dyn().to(device)
    opt=torch.optim.Adam(model.parameters(),lr=2e-3)
    for step in range(steps):
        z=torch.randn(48,2,device=device); th=(torch.rand(48,1,device=device)*2-1)*math.pi
        y=target(z,th)
        pred=model(z,th); loss=((pred-y)**2).mean()
        if penalized:
            lp,_=lyap_penalty(model,z,th); loss=loss+0.8*lp
        opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(),5); opt.step()
    with torch.no_grad():
        z=torch.randn(256,2,device=device); th=(torch.rand(256,1,device=device)*2-1)*math.pi
        mse=((model(z,th)-target(z,th))**2).mean().item()
    # Need jacobians with gradients enabled, no_grad cannot be used.
    z=torch.randn(128,2,device=device); th=(torch.rand(128,1,device=device)*2-1)*math.pi
    lp,lam=lyap_penalty(model,z,th)
    # independent same-regime P_{k+1}=P(theta), to isolate local contraction.
    J=jacobian_batch(model,z,th); P=model.metric(th); ev,U=torch.linalg.eigh(P); Pinv=(U*ev.rsqrt().unsqueeze(-2))@U.transpose(1,2)
    local=torch.linalg.eigvalsh(Pinv@(J.transpose(1,2)@P@J)@Pinv)[:,-1]
    return {'mse':mse,'penalty_batch':lp.item(),'fraction_violating':(lam>0).float().mean().item(),
            'mean_generalized_lambda':lam.mean().item(),'max_generalized_lambda':lam.max().item(),
            'local_contraction_ratio_mean':local.mean().item(),'local_ratio_max':local.max().item()}

if __name__=='__main__':
    out={'exact_check':exact_check(),'baseline':run(False),'lyapunov':run(True)}
    print(json.dumps(out,indent=2))
    with open('results.json','w') as f: json.dump(out,f,indent=2)