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

SEED = 2975
random.seed(SEED)
np.random.seed(SEED)
torch.manual_seed(SEED)

device = "cuda" if torch.cuda.is_available() else "cpu"
try:
    if device == "cuda":
        torch.cuda.empty_cache()
except Exception:
    device = "cpu"
torch.set_default_dtype(torch.float64)

# A residual neural block, z+eta*r(z,u), with an exact state Jacobian.
class ResidualBlock(nn.Module):
    def __init__(self, d=2, hidden=16, eta=0.25, constrained=False, alpha=0.20):
        super().__init__()
        self.d, self.eta, self.constrained, self.alpha = d, eta, constrained, alpha
        self.net = nn.Sequential(nn.Linear(2*d, hidden), nn.Tanh(), nn.Linear(hidden, d))
        if constrained:
            # This is an explicit safe residual parameterization: the learned
            # nonlinear correction is bounded, while the state Jacobian is
            # damped by a fixed contractive skip coefficient.
            self.skip = math.sqrt(1.0 - alpha)
        else:
            self.skip = 1.0

    def forward(self, z, u):
        correction = self.eta * self.net(torch.cat([z, u], dim=-1))
        return self.skip * z + correction

def jacobian_z(model, z, u):
    z = z.detach().requires_grad_(True)
    y = model(z[None], u[None])[0]
    rows = []
    for i in range(y.numel()):
        rows.append(torch.autograd.grad(y[i], z, retain_graph=True)[0])
    return torch.stack(rows)

def max_eig_tensor(a):
    return torch.linalg.eigvalsh((a+a.T)/2)[-1]

def max_eig_symmetric(a):
    return float(max_eig_tensor(a).detach().cpu())

def train(model, constrained, epochs=350):
    # Same tiny regression task for both blocks. It intentionally includes a
    # mildly expansive target, making the robustness/accuracy tradeoff visible.
    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 + 0.25*u + 0.04*torch.sin(2*z)
    opt = torch.optim.Adam(model.parameters(), lr=3e-3)
    for epoch in range(epochs):
        opt.zero_grad()
        pred = model(z, u)
        task = ((pred-target)**2).mean()
        # Exact per-sample Jacobian penalty, corresponding to
        # M^T P M-P <= -alpha P with P=I.
        viol = []
        for k in range(0, len(z), 8):
            M = jacobian_z(model, z[k], u[k])
            q = max_eig_tensor(M.T @ M - (1-model.alpha)*torch.eye(2, device=device))
            viol.append(torch.nn.functional.softplus(q + 0.02)**2)
        penalty = torch.stack(viol).mean()
        loss = task + (2.0*penalty if constrained else 0.0)
        loss.backward()
        opt.step()
    return float(task.detach().cpu()), float(penalty.detach().cpu())

def evaluate(model, alpha):
    g = torch.Generator(device=device).manual_seed(SEED+9)
    points = torch.rand(64, 2, generator=g, device=device)*2-1
    inputs = torch.rand(64, 2, generator=g, device=device)*2-1
    eigs=[]
    for z,u in zip(points, inputs):
        M=jacobian_z(model,z,u)
        eigs.append(max_eig_symmetric(M.T@M-(1-alpha)*torch.eye(2,device=device)))
    # Identical-input trajectories from paired hidden states.
    z1 = torch.tensor([0.65,-0.45], device=device)
    delta = torch.tensor([1e-3,-1.2e-3], device=device)
    z2 = z1 + delta
    u = torch.tensor([0.2,-0.1], device=device)
    ratios=[]
    for _ in range(20):
        old=float(torch.linalg.vector_norm(z2-z1).detach().cpu())
        z1,z2=model(z1[None],u[None])[0],model(z2[None],u[None])[0]
        ratios.append(float((torch.linalg.vector_norm(z2-z1)/old).detach().cpu()))
    return {
        "max_dense_generalized_violation": float(max(eigs)),
        "mean_dense_generalized_violation": float(np.mean(eigs)),
        "max_trajectory_ratio": float(max(ratios)),
        "mean_trajectory_ratio": float(np.mean(ratios)),
        "final_distance_ratio": float(np.prod(ratios)),
        "ratios": ratios,
    }

def direct_math_check():
    # M=0.8I exactly satisfies M^T M-I <= -0.36I (alpha=.30), while
    # M=1.1I violates it and distances grow by 1.1 each step.
    alpha=.30
    good=torch.eye(2)*0.8
    bad=torch.eye(2)*1.1
    good_margin=max_eig_symmetric(good.T@good-(1-alpha)*torch.eye(2))
    bad_margin=max_eig_symmetric(bad.T@bad-(1-alpha)*torch.eye(2))
    dgood=1.0; dbad=1.0
    for _ in range(12): dgood*=.8; dbad*=1.1
    return {"alpha":alpha,"contractive_matrix_largest_eigenvalue":good_margin,
            "expansive_matrix_largest_eigenvalue":bad_margin,
            "contractive_distance_after_12":dgood,
            "expansive_distance_after_12":dbad}

def main():
    global device
    results={"seed":SEED,"device":device,"direct_math_check":direct_math_check()}
    for name, constrained in [("baseline",False),("idea",True)]:
        try:
            model=ResidualBlock(constrained=constrained).to(device)
            task, penalty=train(model, constrained)
            results[name]={"train_task_mse":task,"train_penalty":penalty,**evaluate(model, .20)}
        except Exception as e:
            if device == "cuda":
                device="cpu"; torch.set_default_device("cpu")
                model=ResidualBlock(constrained=constrained)
                task, penalty=train(model, constrained)
                results[name]={"train_task_mse":task,"train_penalty":penalty,**evaluate(model,.20)}
            else: raise
    with open("results.json","w") as f: json.dump(results,f,indent=2)
    print(json.dumps(results,indent=2))

if __name__ == "__main__": main()