import json, math, random, time
from pathlib import Path
import numpy as np
import torch
from torch import nn

SEED=3015
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"
try:
    if device == "cuda": torch.cuda.set_device(0)
except Exception:
    device = "cpu"
DT=torch.float32
NPER=12
ANGLE=2*math.pi/NPER
ROT=torch.tensor([[math.cos(ANGLE),-math.sin(ANGLE)],[math.sin(ANGLE),math.cos(ANGLE)]],dtype=DT)
TRUE_A=.86

def true_map(z):
    # radial contraction toward the unit circle, followed by a fixed rotation
    r=torch.sqrt((z*z).sum(-1,keepdim=True)+1e-9)
    rn=1+TRUE_A*(r-1)
    return (rn/r)*z @ ROT.to(z.device).T

def make_data(ntraj=18, length=90, noise=.035):
    xs=[]; ys=[]
    for _ in range(ntraj):
        z=torch.tensor([1+np.random.uniform(-.18,.18),0.],dtype=DT)
        phase=np.random.uniform(0,2*math.pi)
        z=torch.tensor([math.cos(phase),math.sin(phase)],dtype=DT)*float(np.random.uniform(.82,1.18))
        for k in range(length):
            xn=z+noise*torch.randn(2)
            zn=true_map(z)
            yn=zn+noise*torch.randn(2)
            xs.append(xn); ys.append(yn); z=zn
    return torch.stack(xs).to(device),torch.stack(ys).to(device)

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

def monodromy(model,z):
    # Exact autodiff Jacobian of the N-step return map at the section point.
    zz=z
    for _ in range(NPER): zz=model(zz)
    rows=[]
    for j in range(2): rows.append(torch.autograd.grad(zz[0,j],z,create_graph=True,retain_graph=True)[0][0])
    return torch.stack(rows),zz

def train(kind, x, y, steps=1800):
    model=Dynamics().to(device)
    opt=torch.optim.Adam(model.parameters(),lr=2e-3)
    rng=np.random.default_rng(SEED+ (0 if kind=='baseline' else 1))
    trace=[]
    for step in range(steps):
        ix=torch.tensor(rng.integers(0,len(x),128),device=device)
        pred=model(x[ix]); loss=((pred-y[ix])**2).mean()
        if kind=='idea':
            # Section h(z)=y=0, with a known 2pi symmetry return (alpha=2pi).
            # Start exactly on the section and require P(z)=z.
            z0=torch.tensor([[1.,0.]],device=device,requires_grad=True)
            J,ret=monodromy(model,z0)
            retloss=((ret-z0)**2).mean()
            eig=torch.linalg.eigvals(J)
            rho=torch.abs(eig).max()
            floq=torch.relu(rho-(1-.03))**2
            loss=loss+1.2*retloss+0.25*floq
        opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(),5.0); opt.step()
        if step in (0,steps//2,steps-1): trace.append(float(loss.detach().cpu()))
    with torch.no_grad():
        z0=torch.tensor([[1.,0.]],device=device)
    z0.requires_grad_(True)
    J,ret=monodromy(model,z0)
    rho=float(torch.abs(torch.linalg.eigvals(J)).max().detach().cpu())
    return model, rho, float(((ret-z0)**2).mean().detach().cpu()), trace

def rollout_error(model, nperiods=15):
    # Compare a perturbed radius trajectory with the true dynamics; report error at event returns.
    z=torch.tensor([[.90,0.]],device=device)
    truth=z.clone(); errs=[]
    with torch.no_grad():
        for _ in range(nperiods):
            for _ in range(NPER): z=model(z); truth=true_map(truth)
            errs.append(float(torch.linalg.norm(z-truth).cpu()))
    return errs

def main():
    t=time.time(); x,y=make_data()
    results={'device':device,'seed':SEED,'period_steps':NPER,'true_radial_floquet':TRUE_A}
    alltr={}
    for kind in ('baseline','idea'):
        model,rho,ret,trace=train(kind,x,y)
        errs=rollout_error(model)
        results[kind]={'rho':rho,'return_mse':ret,'event_errors':errs,'loss_trace':trace}
        alltr[kind]=model.state_dict()
    # Independent core math sanity check: finite perturbation ratios for the analytic map.
    z=torch.tensor([[1.001,0.]],dtype=DT); q=torch.tensor([[1.,0.]],dtype=DT)
    ratios=[]
    for n in range(1,6):
        for _ in range(NPER): z=true_map(z); q=true_map(q)
        ratios.append(float((torch.linalg.norm(z-q)/.001).item()))
    results['analytic_perturbation_ratios']=ratios
    results['runtime_sec']=time.time()-t
    Path('results.json').write_text(json.dumps(results,indent=2))
    print(json.dumps(results,indent=2))

if __name__=='__main__': main()