import json, math
import numpy as np

SEED=1349
alpha,beta=4.0,0.2

def drift(x,g): return alpha*x*x+g-x**3-beta*x
def ddrift(x): return 2*alpha*x-3*x*x-beta

def extrema(): return np.sort(np.roots([3.,-2*alpha,beta]).real)

def fold_bounds():
    e=extrema(); q=lambda x: alpha*x*x-x**3-beta*x
    # gamma values at which a pair of fixed points coalesces
    return float(-q(e[1])),float(-q(e[0])),e

def roots(g):
    z=np.roots([1.,-alpha,beta,-g])
    return np.sort(z[np.abs(z.imag)<1e-7].real)

def integrate(gfun,omega,dt=0.004,cycles=6,noise=0.,x0=None):
    T=2*np.pi/omega; n=min(max(int(cycles*T/dt),4000),300000)
    dt=cycles*T/n; x=3.9 if x0 is None else x0
    rng=np.random.default_rng(SEED+int(1000*omega))
    xs=[]; gs=[]
    for k in range(n):
        t=k*dt
        x=max(0.,x+dt*drift(x,gfun(t))+math.sqrt(2*noise*dt)*rng.normal())
        if k>=n//2: xs.append(x); gs.append(gfun(t+dt))
    xs,gs=np.asarray(xs),np.asarray(gs)
    return float(abs(np.sum(.5*(xs[1:]+xs[:-1])*(gs[1:]-gs[:-1])))),float(dt)

def frequency_sweep(g0,amp):
    # Estimate tau on the high stable branch at the center forcing.
    rr=roots(g0); stable=rr[ddrift(rr)<0]
    tau=float(-1/ddrift(stable[-1]))
    omegas=np.logspace(math.log10(.03/tau),math.log10(30/tau),13)
    rows=[]
    for w in omegas:
        a,dt=integrate(lambda t:g0+amp*math.sin(w*t),w)
        rows.append({'omega':float(w),'omega_tau':float(w*tau),'area':a})
    return tau,rows

def classification():
    rng=np.random.default_rng(SEED); N,L=320,40
    X=np.zeros((N,L)); y=np.arange(N)%2
    for i in range(N): X[i]=(-5.0 if y[i]==0 else -3.9)+.08*rng.normal(size=L)
    tr=np.arange(240); te=np.arange(240,320)
    def run(kind):
        st=[]
        for seq in X:
            if kind=='cell':
                x=3.9
                for u in seq:
                    for _ in range(4): x=max(0.,x+.05*drift(x,u))
            else:
                x=0.
                for u in seq: x=.9*x+.1*u
            st.append(x)
        st=np.asarray(st); th=(st[tr[y[tr]==0]].mean()+st[tr[y[tr]==1]].mean())/2
        return float(np.mean((st[te]>th)==y[te])),float(st[y==0].mean()),float(st[y==1].mean())
    return run('linear'),run('cell')

def main():
    lo,hi,e=fold_bounds(); gmid=(lo+hi)/2
    rr=roots(gmid)
    math_check={'extrema':e.tolist(),'predicted_fold_interval':[lo,hi],
      'midpoint_gamma':gmid,'all_roots':rr.tolist(),
      'derivatives':[float(ddrift(x)) for x in rr],
      'stable_root_count':int(np.sum(ddrift(rr)<0))}
    # Sweep spans both folds, so deterministic state switching is possible.
    tau,freq=frequency_sweep(gmid,(hi-lo)*.56)
    wstar=1/tau
    noises=[0.,1e-4,5e-4,1e-3,5e-3,2e-2]
    noise=[]
    for d in noises:
        a,_=integrate(lambda t:gmid+(hi-lo)*.56*math.sin(wstar*t),wstar,noise=d)
        noise.append([d,a])
    result={'parameters':{'alpha':alpha,'beta':beta,'seed':SEED},'math_check':math_check,
      'tau_high':tau,'frequency_sweep':freq,'noise_sweep':noise,
      'classification':{'linear':classification()[0],'cell':classification()[1]}}
    with open('results.json','w') as f: json.dump(result,f,indent=2)
    print(json.dumps(result,indent=2))
if __name__=='__main__': main()