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

SEED=7
rng=np.random.default_rng(SEED)

# ---------- Core phase controller ----------
def control(theta,k):
    z=np.exp(1j*theta); r=z.mean()
    # formula in the prompt
    return (2*k/len(theta))*np.imag(z*np.conj(r))

def grad(theta):
    return -control(theta,1.0)

def phase_run(theta0,k=1.,dt=0.002,T=4.,delta=None,omega=None):
    th=theta0.copy(); n=int(T/dt); V=[]; us=[]; events=[]; held=np.zeros(len(th)); last=0
    if omega is None: omega=np.zeros(len(th))
    for q in range(n+1):
        ue=control(th,k)
        if delta is None:
            held=ue
        elif q==0 or np.linalg.norm(ue-held)>=delta:
            held=ue.copy(); events.append(q*dt)
        z=np.exp(1j*th); V.append(abs(z.mean())**2); us.append(ue.copy())
        if q<n: th += dt*(held+omega)
    return np.asarray(V),np.asarray(us),np.asarray(events),th

def math_checks():
    N=24; theta=rng.normal(0,.22,N) # synchronized perturbation, clear descent
    dt=.0002; T=.25
    rows=[]
    for k in [.25,.5,1.,2.]:
        V,u,ev,_=phase_run(theta,k,dt,T,None)
        # finite difference at initial point compared with exact -k||grad||2
        observed=(V[1]-V[0])/dt
        predicted=-k*np.sum(grad(theta)**2)
        rows.append({'k':k,'initial_Vdot_observed':float(observed),'predicted':float(predicted),'ratio':float(observed/predicted)})
    # event tolerance: error and event rate versus delta
    evrows=[]
    theta2=rng.normal(0,.3,N); omega=np.linspace(-1.0,1.0,N)
    # finite-difference Lipschitz estimate for the exact control under uncontrolled drift
    _,ud,_,_=phase_run(theta2,1.,.001,10.,None,omega)
    M=float(np.percentile(np.linalg.norm(np.diff(ud,axis=0)/.001,axis=1),99))
    for d in [.005,.02,.08,.2]:
        V,u,ev,_=phase_run(theta2,1.,.001,10.,d,omega)
        dw=np.diff(ev)
        mindwell=float(dw.min()) if len(dw) else None
        # held/exact discrepancy reconstructed from held intervals, max just-before event
        # use trigger sampling; report event count and final V
        evrows.append({'delta':d,'events':int(len(ev)),'rate':float(len(ev)/2),'final_V':float(V[-1]),'min_dwell':mindwell,'predicted_dwell_lower_bound':float(d/M)})
    return rows,evrows

# ---------- tiny neural experiment ----------
def train_model(mode, seed=11, epochs=18, force_cpu=False):
    import torch
    torch.manual_seed(seed); np.random.seed(seed)
    device='cpu' if force_cpu else ('cuda' if torch.cuda.is_available() else 'cpu')
    try:
        dev=torch.device(device)
        # sequence label is sign of a noisy temporal sum; common input encourages synchrony
        ntr,nte,L=640,256,24
        g=torch.Generator().manual_seed(seed)
        X=torch.randn(ntr+nte,L,1,device=dev)
        labels=(X.sum(1)[:,0]>0).long()
        # fixed split, same data for all modes
        tr=slice(0,ntr); te=slice(ntr,None)
        N=16; H=2*N; dt=.15; k=1.2; delta=.035
        W=torch.randn(H,H,device=dev)*.22
        U=torch.randn(1,H,device=dev)*.35
        b=torch.zeros(H,device=dev)
        W.requires_grad_(); U.requires_grad_(); b.requires_grad_()
        out=torch.randn(H,2,device=dev)*.15; out.requires_grad_()
        opt=torch.optim.Adam([W,U,b,out],lr=.018)
        event_total=0; steps_total=0; amp_drift=[]
        for ep in range(epochs):
            perm=torch.randperm(ntr,device=dev)
            for st in range(0,ntr,64):
                ids=perm[st:st+64]; xb=X[ids]; yb=labels[ids]; B=xb.shape[0]
                h=torch.zeros(B,H,device=dev); held=torch.zeros(B,N,device=dev)
                for t in range(L):
                    base=torch.tanh(h@W + xb[:,t]@U + b)
                    pairs=base.view(B,N,2)
                    z=pairs[...,0]+1j*pairs[...,1]; z=z/(torch.abs(z)+1e-6); r=z.mean(1,keepdim=True)
                    ue=(2*k/N)*torch.imag(z*torch.conj(r))
                    if mode=='continuous': held=ue
                    elif mode=='event':
                        trigger=torch.linalg.vector_norm(ue-held,dim=1)>=delta
                        held=torch.where(trigger[:,None],ue,held); event_total+=int(trigger.sum())
                    # exact tangent rotation, preserving each pair norm
                    ang=dt*held; x,yy=pairs[...,0],pairs[...,1]
                    rot=torch.stack([x*torch.cos(ang)-yy*torch.sin(ang),x*torch.sin(ang)+yy*torch.cos(ang)],-1)
                    h=rot.reshape(B,H); steps_total+=B
                logits=h@out; loss=torch.nn.functional.cross_entropy(logits,yb)
                opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_([W,U,b,out],1.); opt.step()
        with torch.no_grad():
            h=torch.zeros(nte,H,device=dev); held=torch.zeros(nte,N,device=dev); norm0=None
            for t in range(L):
                pairs=torch.tanh(h@W+X[te,t]@U+b).view(nte,N,2)
                if norm0 is None: norm0=torch.linalg.vector_norm(pairs,dim=-1)
                z=pairs[...,0]+1j*pairs[...,1]; z=z/(torch.abs(z)+1e-6); r=z.mean(1,keepdim=True); ue=(2*k/N)*torch.imag(z*torch.conj(r))
                if mode=='continuous': held=ue
                elif mode=='event': held=torch.where((torch.linalg.vector_norm(ue-held,dim=1)>=delta)[:,None],ue,held)
                ang=dt*held; x,yy=pairs[...,0],pairs[...,1]
                h=torch.stack([x*torch.cos(ang)-yy*torch.sin(ang),x*torch.sin(ang)+yy*torch.cos(ang)],-1).reshape(nte,H)
            acc=(h@out).argmax(1).eq(labels[te]).float().mean().item()
            # compare base pair norms before/after rotation is exact up to float error
            normerr=float((torch.linalg.vector_norm(h.view(nte,N,2),dim=-1)-torch.linalg.vector_norm(pairs,dim=-1)).abs().mean())
        return {'accuracy':acc,'events_per_sequence':event_total/max(1,ntr*epochs),'rotation_norm_error':normerr,'device':str(dev)}
    except Exception as e:
        if device=='cuda':
            return train_model(mode, seed, epochs, force_cpu=True)
        return {'error':repr(e)}

def main():
    mathrows,eventrows=math_checks()
    results={}
    for mode in ['none','continuous','event']:
        results[mode]=train_model(mode)
    out={'seed':SEED,'math_gain_sweep':mathrows,'event_delta_sweep':eventrows,'neural_results':results}
    Path('results.json').write_text(json.dumps(out,indent=2))
    print(json.dumps(out,indent=2))
if __name__=='__main__': main()