import json, math, random
from pathlib import Path
import numpy as np
from scipy.special import lambertw
import torch
from torch import nn

SEED = 1082
random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)

# x'(t)=a x(t)+b x(t-tau), with characteristic roots
# lambda = a + W_k(b*tau*exp(-a*tau))/tau.
def roots_discrete(a, b, tau, branches=range(-40, 41)):
    if tau == 0:
        return np.array([a+b], dtype=complex)
    z = b*tau*np.exp(-a*tau)
    return np.array([a + lambertw(z, k)/tau for k in branches], dtype=complex)

def rightmost(a, b, tau):
    rr = roots_discrete(a,b,tau)
    return rr[np.argmax(rr.real)]

def hopf_prediction(a,b,k=0):
    w = math.sqrt(b*b-a*a)
    # Direct transcription of the supplied atan2 formula, normalized to tau>0.
    phase = math.atan2(-w/b, -a/b) + 2*math.pi*k
    while phase <= 0: phase += 2*math.pi
    return phase/w, w

def math_check():
    # Parameter sweep: each case has |b|>|a| and therefore a Hopf candidate.
    cases=[(-1.0,-2.0),(-0.5,-1.5),(-1.5,-2.5),(-0.8,-1.8)]
    rows=[]
    for a,b in cases:
        tau_pred,w_pred=hopf_prediction(a,b)
        taus=np.linspace(0.02, tau_pred*2.0, 700)
        rs=np.array([rightmost(a,b,t).real for t in taus])
        ix=np.where(np.sign(rs[:-1]) != np.sign(rs[1:]))[0][0]
        tc=taus[ix] - rs[ix]*(taus[ix+1]-taus[ix])/(rs[ix+1]-rs[ix])
        rc=rightmost(a,b,tc)
        eps=max(.001,tau_pred*0.005)
        slope=(rightmost(a,b,tc+eps).real-rightmost(a,b,tc-eps).real)/(2*eps)
        freq=float(abs(rc.imag))
        below=rightmost(a,b,tau_pred*.8).real
        above=rightmost(a,b,tau_pred*1.2).real
        rows.append({"a":a,"b":b,"predicted_tau_c":tau_pred,
          "observed_tau_c":float(tc),"tau_error_pct":float(abs(tc-tau_pred)/tau_pred*100),
          "predicted_omega":w_pred,"observed_omega":freq,
          "omega_error_pct":float(abs(freq-w_pred)/w_pred*100),
          "dr_dtau_observed":float(slope),"stable_side_r":float(below),
          "unstable_side_r":float(above),"confirmed":bool(abs(tc-tau_pred)/tau_pred<.02 and
          abs(freq-w_pred)/w_pred<.02 and slope>0 and below<0<above)})
    return {"cases":rows,"all_predictions_confirmed":all(x["confirmed"] for x in rows),
            "max_tau_error_pct":max(x["tau_error_pct"] for x in rows),
            "max_omega_error_pct":max(x["omega_error_pct"] for x in rows)}

class DelayedRNN(nn.Module):
    def __init__(self, d, hidden=24):
        super().__init__(); self.d=d; self.hidden=hidden
        self.inp=nn.Linear(1,hidden); self.rec=nn.Linear(hidden,hidden,bias=False)
        self.out=nn.Linear(hidden,2)
    def forward(self,x):
        B,T,_=x.shape; hs=[torch.zeros(B,self.hidden,device=x.device) for _ in range(self.d+1)]
        for t in range(T):
            h=torch.tanh(self.inp(x[:,t])+self.rec(hs[-self.d]))
            hs.append(h); hs.pop(0)
        return self.out(hs[-1])

def data(n, T=32):
    # Two frequencies, with random phase and modest noise.
    y=np.random.randint(0,2,n); phase=np.random.rand(n)*2*np.pi
    t=np.arange(T)[None,:]
    f=np.where(y[:,None]==0,.12,.25)
    x=np.sin(2*np.pi*f*t+phase[:,None])+0.08*np.random.randn(n,T)
    return torch.tensor(x[:,:,None],dtype=torch.float32), torch.tensor(y,dtype=torch.long)

def train(scheduled, device):
    torch.manual_seed(SEED+int(scheduled)); np.random.seed(SEED+int(scheduled))
    tr_y=data(768); va_y=data(256)
    tr=(tr_y[0].to(device),tr_y[1].to(device)); va=(va_y[0].to(device),va_y[1].to(device))
    model=DelayedRNN(1,24).to(device); opt=torch.optim.Adam(model.parameters(),lr=3e-3); lossfn=nn.CrossEntropyLoss()
    a,b=-1.,-2.; target=-.10; tau=.25; history=[]
    # In this tiny implementation integer buffer delay is the exposed curriculum variable.
    for step in range(120):
        if scheduled and step%10==0:
            r=rightmost(a,b,tau).real
            # supplied control law, with clipping and a stable warm-up target
            tau=float(np.clip(tau+0.30*(target-r),.05,2.0))
            history.append((step,tau,float(r)))
        d=max(1,min(8,int(round(tau*4))))
        model.d=d
        ix=torch.randint(0,768,(64,),device=device)
        opt.zero_grad(); loss=lossfn(model(tr[0][ix]),tr[1][ix]); loss.backward(); opt.step()
    with torch.no_grad():
        pred=model(va[0]).argmax(1); acc=float((pred==va[1]).float().mean().cpu())
        vl=float(lossfn(model(va[0]),va[1]).cpu())
    return {"val_loss":vl,"accuracy":acc,"final_tau":tau,"final_integer_delay":max(1,min(8,int(round(tau*4)))),"schedule":history}

def main():
    check=math_check()
    device='cuda' if torch.cuda.is_available() else 'cpu'
    try:
        baseline=train(False,device); idea=train(True,device)
    except Exception as e:
        device='cpu'; baseline=train(False,device); idea=train(True,device)
    result={"math_check":check,"device":device,"baseline":baseline,"idea":idea}
    Path('results.json').write_text(json.dumps(result,indent=2))
    print(json.dumps(result,indent=2))
if __name__=='__main__': main()