import json, math, random, time
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)
device = "cuda" if torch.cuda.is_available() else "cpu"
# A conservative fallback: all computations fit easily on CPU or a small GPU.
try:
    if device == "cuda":
        torch.cuda.empty_cache()
except Exception:
    device = "cpu"

dtype = torch.float32
omega = torch.tensor([1.0, math.sqrt(2.0)], device=device, dtype=dtype)

# K(theta) = [cos(theta1), sin(theta1), cos(theta2), sin(theta2)].
# This is a finite Fourier embedding; coefficients are trainable in the experiment.
class FourierTorus(nn.Module):
    def __init__(self):
        super().__init__()
        # columns: cosine/sine coefficients for each phase and state coordinate
        self.a = nn.Parameter(torch.tensor([[1.,0.],[0.,1.],[0.,0.],[0.,0.]], device=device))
        self.b = nn.Parameter(torch.tensor([[0.,0.],[1.,0.],[0.,0.],[0.,1.]], device=device))
    def forward(self, theta):
        # theta [N,2], output [N,4]
        return torch.cos(theta) @ self.a.T + torch.sin(theta) @ self.b.T
    def tangent(self, theta):
        # dK/dtheta_j * omega_j, exact Fourier differentiation
        return (-torch.sin(theta) * omega) @ self.a.T + (torch.cos(theta) * omega) @ self.b.T

class Field(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(4, 48), nn.Tanh(), nn.Linear(48, 48), nn.Tanh(), nn.Linear(48,4))
    def forward(self, x): return self.net(x)

def true_field(x):
    y = torch.zeros_like(x)
    y[:,0] = -omega[0]*x[:,1]; y[:,1] = omega[0]*x[:,0]
    y[:,2] = -omega[1]*x[:,3]; y[:,3] = omega[1]*x[:,2]
    return y

def rk4(model, x, dt):
    k1=model(x); k2=model(x+dt*k1/2); k3=model(x+dt*k2/2); k4=model(x+dt*k3)
    return x + dt*(k1+2*k2+2*k3+k4)/6

def make_data(n=192, noise=0.025):
    # Short observations: only a fraction of the long quasiperiodic beat.
    t = torch.rand(n, device=device)*7.0
    phases = torch.stack([0.4 + t, 1.1 + math.sqrt(2.0)*t], dim=1)
    x = torch.stack([torch.cos(phases[:,0]), torch.sin(phases[:,0]), torch.cos(phases[:,1]), torch.sin(phases[:,1])], dim=1)
    y = true_field(x) + noise*torch.randn_like(x)
    return x, y

def train(use_torus, seed=SEED, steps=1400):
    torch.manual_seed(seed)
    x, y = make_data()
    field, torus = Field().to(device), FourierTorus().to(device)
    opt = torch.optim.Adam(list(field.parameters()) + (list(torus.parameters()) if use_torus else []), lr=3e-3)
    L=12
    grid = torch.stack(torch.meshgrid(torch.arange(L, device=device)*2*math.pi/L,
                                      torch.arange(L, device=device)*2*math.pi/L, indexing='ij'), dim=-1).reshape(-1,2)
    for step in range(steps):
        opt.zero_grad()
        pred = field(x)
        loss = ((pred-y)**2).mean()
        if use_torus:
            K = torus(grid); T = torus.tangent(grid)
            colloc = ((field(K)-T)**2).mean()
            # Gauge: preserve the nonconstant Fourier embedding amplitudes and avoid collapse.
            gauge = (torus.a.pow(2).sum(dim=0)-1).pow(2).mean() + (torus.b.pow(2).sum(dim=0)-1).pow(2).mean()
            loss = loss + 0.8*colloc + 0.03*gauge
        loss.backward(); opt.step()
    # Evaluate on a long trajectory with exact starting point.
    t = torch.linspace(0, 80, 1601, device=device)
    ph = torch.stack([0.4+t, 1.1+math.sqrt(2.0)*t], dim=1)
    truth = torch.stack([torch.cos(ph[:,0]),torch.sin(ph[:,0]),torch.cos(ph[:,1]),torch.sin(ph[:,1])], dim=1)
    z=truth[0:1].clone(); traj=[z]
    with torch.no_grad():
        for _ in range(1600):
            z=rk4(field,z,0.05); traj.append(z)
        predtraj=torch.cat(traj)
        err=torch.sqrt(((predtraj-truth)**2).sum(dim=1))
        radii=torch.stack([torch.sqrt(predtraj[:,0]**2+predtraj[:,1]**2), torch.sqrt(predtraj[:,2]**2+predtraj[:,3]**2)],dim=1)
        one_step=((field(x)-true_field(x))**2).mean().sqrt().item()
        # Collocation residual on the same fixed grid.
        cres=((field(torus(grid))-torus.tangent(grid))**2).mean().sqrt().item() if use_torus else float('nan')
    return {'long_rmse': float(torch.sqrt((predtraj-truth).pow(2).mean()).item()),
            'final_error': float(err[-1].item()), 'max_radius_dev': float((radii-1).abs().max().item()),
            'one_step_field_rmse':one_step, 'collocation_rmse':cres}

def derivative_check():
    torch.manual_seed(3); tor=FourierTorus().to(device)
    th=torch.rand(11,2,device=device)*2*math.pi
    eps=1e-3
    direction=torch.tensor([0.37,-0.52],device=device)
    fd=(tor(th+eps*direction)-tor(th-eps*direction))/(2*eps)
    # Check the actual Fourier derivative independently, using arbitrary direction.
    actual=(-torch.sin(th)*direction)@tor.a.T + (torch.cos(th)*direction)@tor.b.T
    return float((fd-actual).abs().max().item()), float((tor.tangent(th)-((-torch.sin(th)*omega)@tor.a.T+(torch.cos(th)*omega)@tor.b.T)).abs().max().item())

if __name__ == '__main__':
    t0=time.time(); fd_err, tangent_err=derivative_check()
    base=[train(False, SEED+i) for i in range(3)]
    idea=[train(True, SEED+i) for i in range(3)]
    def avg(rows):
        return {k: float(np.mean([r[k] for r in rows])) for k in rows[0]}
    out={'device':device,'derivative_fd_max_abs':fd_err,'analytic_tangent_self_check':tangent_err,
         'baseline':avg(base),'idea':avg(idea),'baseline_runs':base,'idea_runs':idea,
         'seconds':time.time()-t0}
    print(json.dumps(out, indent=2)); open('results.json','w').write(json.dumps(out, indent=2, allow_nan=True))