import json, math, os
import numpy as np

# Bifurcation-aware adaptive RK4 toy experiment.
# The scalar normal form is dz/dt = mu + z^2, with passage measured from z=-1 to z=+1.

def rhs(z, mu):
    return mu + z*z

def rk4_step(z, h, mu):
    k1 = rhs(z, mu)
    k2 = rhs(z + .5*h*k1, mu)
    k3 = rhs(z + .5*h*k2, mu)
    k4 = rhs(z + h*k3, mu)
    return z + h*(k1 + 2*k2 + 2*k3 + k4)/6

def exact_residence(mu):
    # Integral dz/(mu+z^2), for positive mu, between -1 and +1.
    q = math.sqrt(mu)
    return 2.0*math.atan(1.0/q)/q

def controller_h(mu, hmin=5e-4, hmax=5e-2, mu0=1e-2, delta=1e-12):
    # Formula in the proposal, with mu_hat clipped to its positive ghost regime.
    return min(hmax, max(hmin, hmax*math.sqrt((max(mu, 0.0)+delta)/mu0)))

def passage(mu, mode, fixed_h=2e-3):
    z, t, n = -1.0, 0.0, 0
    while z < 1.0 and n < 100000000:
        h = fixed_h if mode == 'fixed' else controller_h(mu)
        # Do not step far beyond the event; this also makes errors comparable.
        h = min(h, (1.0-z)/max(rhs(z, mu), 1e-30))
        z = rk4_step(z, h, mu)
        t += h; n += 1
    return t, n, z

def log_slope(x, y):
    return float(np.polyfit(np.log(x), np.log(y), 1)[0])

def slow_passage(eps):
    # mu=eps*t. Start at mu=-0.1 on the attracting quasi-static branch,
    # then measure time from the fold (t=0) until z crosses zero.
    t = -0.1/eps
    z = -math.sqrt(.1)
    # A modest step relative to the universal eps^(1/3) time scale.
    h = min(0.02, 0.03*eps**(-1/3))
    n = 0
    while z < 0.0 and n < 20000000:
        # RK4 with time-varying mu.
        def f(tt, zz): return eps*tt + zz*zz
        k1=f(t,z); k2=f(t+h/2,z+h*k1/2); k3=f(t+h/2,z+h*k2/2); k4=f(t+h,z+h*k3)
        zn=z+h*(k1+2*k2+2*k3+k4)/6
        if zn >= 0.0:
            # linear interpolation of the crossing, sufficient for scaling.
            frac=(0.0-z)/(zn-z)
            t=t+frac*h
            break
        z=zn; t += h; n += 1
    return t, eps*t, n

def main():
    np.random.seed(0)
    mus=np.logspace(-4, -0.5, 8)
    residence=[]; fixed=[]; adaptive=[]
    for mu in mus:
        exact=exact_residence(mu)
        tf,nf,_=passage(mu,'fixed')
        ta,na,_=passage(mu,'adaptive')
        residence.append(exact)
        fixed.append({'time':tf,'n':nf,'relerr':abs(tf-exact)/exact})
        adaptive.append({'time':ta,'n':na,'relerr':abs(ta-exact)/exact,'h':controller_h(mu)})
    # Prediction 1: ghost residence has slope -1/2 at small positive mu.
    small=mus[:5]
    residence_slope=log_slope(small, np.array(residence[:5]))
    # Prediction 2: controller h ~ sqrt(mu), hence controller allocations over the
    # same physical ghost crossing scale approximately mu^-1 (T~mu^-1/2 and 1/h~mu^-1/2).
    alloc_slope=log_slope(mus[:5], np.array([x['n'] for x in adaptive[:5]],float))
    # Prediction 3: slow saddle-node passage time scales eps^-1/3.
    eps=np.logspace(-4,-1,7)
    slow=[slow_passage(e) for e in eps]
    slow_times=np.array([x[0] for x in slow])
    slow_slope=log_slope(eps, slow_times)
    out={
      'predictions':{
        'ghost_residence_slope_pred':-0.5,'ghost_residence_slope_obs':residence_slope,
        'controller_allocation_slope_pred':-1.0,'controller_allocation_slope_obs':alloc_slope,
        'slow_delay_slope_pred':-1/3,'slow_delay_slope_obs':slow_slope},
      'mu_sweep':[{'mu':float(m),'exact_T':float(residence[i]),'fixed':fixed[i], 'adaptive':adaptive[i],
                   'fixed_over_adaptive_n':fixed[i]['n']/max(adaptive[i]['n'],1)} for i,m in enumerate(mus)],
      'slow_sweep':[{'epsilon':float(e),'delay_time':float(slow[i][0]),'mu_at_crossing':float(slow[i][1]),'steps':slow[i][2]} for i,e in enumerate(eps)],
      'summary':{
        'mean_fixed_relerr':float(np.mean([x['relerr'] for x in fixed])),
        'mean_adaptive_relerr':float(np.mean([x['relerr'] for x in adaptive])),
        'far_mu_eval_saving':float(1-adaptive[-1]['n']/fixed[-1]['n']),
        'near_mu_eval_saving':float(1-adaptive[0]['n']/fixed[0]['n'])}}
    with open('results.json','w') as f: json.dump(out,f,indent=2)
    print(json.dumps(out,indent=2))

if __name__=='__main__': main()