import json
from pathlib import Path
import numpy as np
import math
import torch
from torch import nn

SEED = 17
np.random.seed(SEED)
torch.manual_seed(SEED)
device = "cuda" if torch.cuda.is_available() else "cpu"
try:
    if device == "cuda": torch.cuda.set_device(0)
except Exception:
    device = "cpu"

# Exact underdamped oscillator labels, vectorized.
# q'' + .7 q' + 1.5 q = 0, q(0)=p, q'(0)=v.
def true_flow(x, t=1.0):
    p, v = x[:, 0], x[:, 1]
    alpha = 0.35
    omega = (1.5 - alpha * alpha) ** 0.5
    e = math.exp(-alpha * t)
    c, s = math.cos(omega * t), math.sin(omega * t)
    q = e * (p * c + (v + alpha * p) * s / omega)
    qv = e * (v * c - (alpha * v + 1.5 * p) * s / omega)
    return torch.stack((q, qv), dim=1)

class Field(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(2, 24), nn.Tanh(), nn.Linear(24, 24), nn.Tanh(), nn.Linear(24, 2))
    def forward(self, z): return self.net(z)

def rollout(model, z, factors=None, steps=16, T=1.0):
    dt = T / steps
    for i in range(steps):
        fac = 1.0 if factors is None else factors[:, i:i+1]
        z = z + dt * fac * model(z)
    return z

def train(use_ss):
    torch.manual_seed(SEED)
    model = Field().to(device)
    opt = torch.optim.Adam(model.parameters(), lr=4e-3)
    for _ in range(180):
        x = torch.rand(64, 2, device=device) * 2 - 1
        y = true_flow(x)
        nominal = rollout(model, x)
        loss = ((nominal - y) ** 2).mean()
        if use_ss:
            eps = 0.08
            factors = 1 + (2 * torch.rand(64, 16, device=device) - 1) * eps
            pseudo = rollout(model, x, factors)
            # rho is deliberately nonzero as in the proposed hinge loss.
            ss = torch.relu(torch.linalg.vector_norm(pseudo - nominal, dim=1) - .08).pow(2).mean()
            loss = loss + 2.0 * ss
        opt.zero_grad(); loss.backward(); opt.step()
    with torch.no_grad():
        x = torch.rand(256, 2, device=device) * 2 - 1
        y = true_flow(x)
        nominal = rollout(model, x)
        factors = 1 + (2 * torch.rand(256, 16, device=device) - 1) * .08
        pseudo = rollout(model, x, factors)
        return {"clean_rmse": float(((nominal-y)**2).mean().sqrt()),
                "timing_perturbed_rmse": float(((pseudo-y)**2).mean().sqrt()),
                "mean_standard_tracking_error": float(torch.linalg.vector_norm(pseudo-nominal,dim=1).mean())}

if __name__ == "__main__":
    try:
        result = {"device": device, "baseline": train(False), "standard_shadowing": train(True)}
    except Exception as e:
        torch.manual_seed(SEED); device = "cpu"
        result = {"device": "cpu", "error_fallback": repr(e), "baseline": train(False), "standard_shadowing": train(True)}
    Path("neural_ode_results.json").write_text(json.dumps(result, indent=2))
    print(json.dumps(result, indent=2))