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

def cubic_rate(gamma, kappa):
    if kappa <= 0: return 0.0
    lo, hi = 0.0, kappa**(1/3) + math.sqrt(kappa/gamma) + 1.0
    for _ in range(100):
        mid=(lo+hi)/2
        if mid**3+gamma*mid**2 < kappa: lo=mid
        else: hi=mid
    return (lo+hi)/2

def third_order_step(x,v,a,grad,dt,gamma,eps=0.,rng=None):
    a += dt*(-grad-gamma*a)
    if eps: a += math.sqrt(2*gamma*eps*dt)*rng.normal()
    v += dt*a; x += dt*v
    return x,v,a

def measure_growth(kappa,gamma,dt,T=None):
    r=cubic_rate(gamma,kappa)
    if T is None: T=max(60.,20./max(r,1e-5))
    x,v,a=1e-8,0.,0.; n=int(T/dt); ts=[]; ys=[]
    for i in range(n):
        x,v,a=third_order_step(x,v,a,-kappa*x,dt,gamma)
        if i>n*.45 and np.isfinite(x) and 1e-14<abs(x)<1e50:
            ts.append(i*dt); ys.append(math.log(abs(x)))
    if len(ts)<20: return float('nan')
    return float(np.polyfit(ts,ys,1)[0])

def double_well_escape(method,gamma=1.,dt=.01,eps=.08,trials=120,T=35.,seed=4):
    rng=np.random.default_rng(seed); hits=[]
    for _ in range(trials):
        x=-1.;v=0.;a=0.;hit=None
        for j in range(int(T/dt)):
            grad=4*x*(x*x-1)
            if method=='third': x,v,a=third_order_step(x,v,a,grad,dt,gamma,eps,rng)
            elif method=='momentum':
                v=gamma*v-dt*grad+math.sqrt(2*eps*dt)*rng.normal(); x+=v
            else: x-=dt*grad-math.sqrt(2*eps*dt)*rng.normal()
            if x>0: hit=(j+1)*dt; break
        if hit is not None: hits.append(hit)
    return {'hit_fraction':len(hits)/trials,'median_hit':float(np.median(hits)) if hits else None}

def main():
    # Prediction 1: asymptotic saddle growth is the positive cubic root.
    gamma,kappa,dt=1.3,2.7,.001
    pred=cubic_rate(gamma,kappa); obs=measure_growth(kappa,gamma,dt,T=30.)
    p1={'gamma':gamma,'kappa':kappa,'dt':dt,'predicted_rate':pred,'observed_rate':obs,'relative_error':abs(obs-pred)/pred}

    # Prediction 2: weak damping gives r proportional to kappa^(1/3).
    gamma=.03; kappas=np.array([.125,.25,.5,1.,2.,4.])
    pred=np.array([cubic_rate(gamma,k) for k in kappas])
    obs=np.array([measure_growth(k,gamma,.002) for k in kappas])
    p2={'gamma':gamma,'kappas':kappas.tolist(),'predicted_rates':pred.tolist(),'observed_rates':obs.tolist(),
        'predicted_loglog_slope':float(np.polyfit(np.log(kappas),np.log(pred),1)[0]),
        'observed_loglog_slope':float(np.polyfit(np.log(kappas),np.log(obs),1)[0]),
        'relative_rate_errors':(abs(obs-pred)/pred).tolist()}

    # Prediction 3: r is zero at zero negative curvature and increases monotonically with kappa.
    ks=np.array([0.,1e-5,1e-3,.01,.1,1.])
    rows=[]
    for k in ks:
        r=cubic_rate(1.,k); measured=0. if k==0 else measure_growth(k,1.,.01)
        rows.append({'kappa':float(k),'predicted_rate':r,'observed_rate':measured})
    p3={'rows':rows,'predicted_monotone':bool(np.all(np.diff([x['predicted_rate'] for x in rows])>=0)),
        'observed_monotone':bool(np.all(np.diff([x['observed_rate'] for x in rows])>=0))}

    # Prediction 4 (integration heuristic): dt*r=.2 should retain small rate distortion.
    gamma,kappa=1.,4.; r=cubic_rate(gamma,kappa); rows=[]
    for dt in [.01,.05,.1,.15,.2,.3,.5]:
        o=measure_growth(kappa,gamma,dt,T=20.)
        rows.append({'dt':dt,'dt_times_rate':dt*r,'observed_rate':o,'relative_error':abs(o-r)/r})
    p4={'rate':r,'recommended_dt_bound':.2/r,'rows':rows}

    bench={m:double_well_escape(m,seed=10+i) for i,m in enumerate(['third','momentum','overdamped'])}
    result={'prediction_checks':{'rate_at_fixed_params':p1,'curvature_scaling':p2,'curvature_transition':p3,'integration_heuristic':p4},'double_well_escape':bench}
    Path('results.json').write_text(json.dumps(result,indent=2)); print(json.dumps(result,indent=2))
if __name__=='__main__': main()