import json
import numpy as np

# Controlled stationary sweep for rho(x|a)=N(0,1/a), with Phi=a*x^2/2+log(a)/2.
# For adot=v, u(x,a)=-(v/(2a))*x solves div(u)-u*d_x Phi=v*d_a Phi.

def covariance_check(seed=7, n=100000):
    rng=np.random.default_rng(seed); out=[]
    for a in (0.5, 1.0, 2.0):
        x=rng.normal(size=n)/np.sqrt(a)
        dphi=.5*(x*x-1/a)
        cov=np.mean((x*x-np.mean(x*x))*(dphi-np.mean(dphi)))
        out.append((a, float(-cov), float(-1/a**2), float(abs(-cov+1/a**2))))
    return out

def sweep(v, controlled, seed, dt=.01, n=10000, a0=.5, a1=2.0):
    rng=np.random.default_rng(seed); x=rng.normal(size=n)/np.sqrt(a0)
    a=a0; records=[]; steps=int(np.ceil((a1-a0)/(v*dt)))
    for k in range(steps+1):
        if k % max(1,steps//100)==0 or k==steps: records.append((a,np.mean(x*x),1/a))
        if k==steps: break
        rate=a+(v/(2*a) if controlled else 0.)
        decay=np.exp(-rate*dt); variance=(1-decay*decay)/rate
        x=decay*x+np.sqrt(variance)*rng.normal(size=n); a=min(a1,a+v*dt)
    z=np.asarray(records); e=z[1:,1]-z[1:,2]
    return z,float(np.mean(e)),float(np.sqrt(np.mean(e*e)))

def deterministic_error(v, controlled, dt=.01, a0=.5, a1=2.0):
    # Exact expected q=<x^2> recurrence for the frozen-coefficient OU update.
    a=a0; q=1/a; steps=int(np.ceil((a1-a0)/(v*dt))); errors=[]
    for _ in range(steps):
        rate=a+(v/(2*a) if controlled else 0.)
        d=np.exp(-rate*dt); q=d*d*q+(1-d*d)/rate
        a=min(a1,a+v*dt); errors.append(q-1/a)
    return float(np.mean(errors)),float(np.sqrt(np.mean(np.square(errors))))

def main():
    rates=np.array([.01,.02,.04,.08,.16])
    sim=[]; det=[]
    for i,v in enumerate(rates):
        _,su,ru=sweep(v,False,100+i); _,sc,rc=sweep(v,True,200+i)
        mu,ru_det=deterministic_error(v,False); mc,rc_det=deterministic_error(v,True)
        sim.append((v,su,sc,ru,rc)); det.append((v,mu,mc,ru_det,rc_det))
    sim=np.asarray(sim); det=np.asarray(det)
    # RMS deterministic errors test the predicted O(v) versus O(v^2) scaling.
    slope_u=float(np.polyfit(np.log(rates),np.log(det[:,3]),1)[0])
    slope_c=float(np.polyfit(np.log(rates),np.log(det[:,4]),1)[0])
    rng=np.random.default_rng(123); a=1.3; v=.17; x=rng.normal(size=100000)/np.sqrt(a)
    lhs=-v/(2*a)-(-(v/(2*a))*x)*(a*x); rhs=v*.5*(x*x-1/a)
    result={
      'model':'rho=N(0,1/a), A=x^2, a: 0.5 -> 2.0',
      'predictions':[
        'covariance response is -1/a^2',
        'the analytic control has zero continuity residual',
        'uncontrolled tracking error scales O(v), controlled error O(v^2) near the adiabatic regime'
      ],
      'covariance_check':covariance_check(),
      'simulation_rates_v_signed_uncontrolled_controlled_rms_uncontrolled_controlled':sim.tolist(),
      'deterministic_rates_v_mean_uncontrolled_controlled_rms_uncontrolled_controlled':det.tolist(),
      'deterministic_loglog_rms_slope_uncontrolled':slope_u,
      'deterministic_loglog_rms_slope_controlled':slope_c,
      'continuity_residual_rms':float(np.sqrt(np.mean((lhs-rhs)**2))),
      'fastest_rate_controlled_to_uncontrolled_rms_ratio':float(det[-1,4]/det[-1,3])
    }
    with open('results.json','w') as f: json.dump(result,f,indent=2)
    print(json.dumps(result,indent=2))

if __name__=='__main__': main()