import math, random, json
from pathlib import Path
import numpy as np

SEED=7
np.random.seed(SEED); random.seed(SEED)

def metrics(phi,dt):
    w=np.diff(np.unwrap(phi))/dt
    return float(w.mean()),float(w.var()),float(np.sqrt(w.var())/(abs(w.mean())+1e-12))

def math_sweeps():
    omega,dt=1.7,.002
    t=np.arange(0,30*2*np.pi/omega,dt); th=omega*t
    scaling=[]
    for m in [1,2,3,5]:
      for a in [0,.01,.03,.06]:
        _,v,e=metrics(th+a*np.sin(m*th),dt)
        scaling.append({'m':m,'a':a,'observed_var':v,'predicted_var':(omega*a*m)**2/2,'observed_E':e,'predicted_E':abs(a*m)/math.sqrt(2)})
    noise=[]
    for sig in [0,.001,.003,.01,.03]:
      vals=[]
      for r in range(30):
        rng=np.random.default_rng(SEED+r)
        vals.append(metrics(th+rng.normal(0,sig,len(th)),dt)[1])
      pred=2*sig**2/dt**2
      noise.append({'sigma':sig,'observed_var':float(np.mean(vals)),'predicted_var':pred})
    boundary=[]
    for inc in [0.5,1,2,2.8,3.0,3.13,3.15,3.3,4.0]:
      ph=np.arange(1000)*inc
      mw,v,e=metrics(ph,1.0)
      boundary.append({'increment':inc,'estimated_mean':mw,'error':abs(mw-inc),'E':e})
    return scaling,noise,boundary

def gru_experiment():
  try:
    import torch
    import torch.nn as nn
    torch.manual_seed(SEED); device='cuda' if torch.cuda.is_available() else 'cpu'
    dt=.1; omega=1.0; T=45; ntrain=96; ntest=24
    def make(n,seed):
      rng=np.random.default_rng(seed); x=[]
      for i in range(n):
        phase=rng.uniform(0,2*np.pi); amp=1+rng.normal(0,.04)
        tt=np.arange(T+1)*dt
        z=np.stack([amp*np.sin(omega*tt+phase),amp*np.cos(omega*tt+phase)],1)
        z+=rng.normal(0,.025,z.shape); x.append(z)
      return torch.tensor(np.array(x),dtype=torch.float32,device=device)
    train,test=make(ntrain,10),make(ntest,20)
    class Model(nn.Module):
      def __init__(self,clock):
        super().__init__(); self.clock=clock; self.gru=nn.GRU(2,16,batch_first=True); self.out=nn.Linear(16,2)
        self.phase=nn.Linear(16,2) if clock else None
      def forward(self,x):
        h,_=self.gru(x); y=self.out(h); return y,h
    def run(clock, dev):
      torch.manual_seed(SEED); m=Model(clock).to(dev); opt=torch.optim.Adam(m.parameters(),lr=3e-3)
      tr=train.to(dev); te=test.to(dev)
      # teacher-forced one-step prediction; phase head is intentionally auxiliary.
      for ep in range(220):
        inp,tar=tr[:,:-1],tr[:,1:]; y,h=m(inp); loss=((y-tar)**2).mean()
        if clock:
          q=m.phase(h); ph=torch.atan2(q[...,0],q[...,1]); d=torch.atan2(torch.sin(torch.diff(ph,dim=1)),torch.cos(torch.diff(ph,dim=1)))/dt
          # target natural rate is known in this synthetic benchmark; cycle term prevents zero-clock collapse.
          loss=loss+.08*d.var(dim=1).mean()+.08*(d.mean(dim=1)-omega).pow(2).mean()
        opt.zero_grad(); loss.backward(); opt.step()
      with torch.no_grad():
        y,h=m(te[:,:-1]); one=float(((y-te[:,1:])**2).mean().cpu())
        # assess phase consistency of learned phase head; baseline uses geometric atan2.
        if clock: ph=torch.atan2(m.phase(h)[...,0],m.phase(h)[...,1])
        else: ph=torch.atan2(h[...,0],h[...,1])
        ds=torch.atan2(torch.sin(torch.diff(ph,dim=1)),torch.cos(torch.diff(ph,dim=1)))/dt
        return {'one_step_mse':one,'clock_E':float((ds.std(1)/(ds.mean(1).abs()+1e-6)).mean().cpu()),'device':str(dev)}
    try:
      return {'baseline':run(False,device),'clock':run(True,device)}
    except Exception as cuda_error:
      # Required robust fallback after any CUDA/cuDNN allocation or execution error.
      if device != 'cpu':
        return {'baseline':run(False,'cpu'),'clock':run(True,'cpu'),'cuda_error':str(cuda_error)}
      raise
  except Exception as e:
    return {'error':str(e),'fallback':'math-only'}

def main():
  s,n,b=math_sweeps(); out={'scaling':s,'noise':n,'unwrap_boundary':b,'gru':gru_experiment()}
  Path('results.json').write_text(json.dumps(out,indent=2))
  print(json.dumps(out,indent=2))
if __name__=='__main__': main()