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

OUT = Path('results.json')

def roots(lam, kp, ki):
    return np.roots([1.0, kp*lam, ki*lam])

def rk4_step(y, h, lam, kp, ki):
    def f(v):
        th, z = v
        return np.array([-kp*lam*th-ki*lam*z, th])
    k1=f(y); k2=f(y+h*k1/2); k3=f(y+h*k2/2); k4=f(y+h*k3)
    return y+h*(k1+2*k2+2*k3+k4)/6

def continuous(lam, kp, ki, reset=False, dwell=0.0, T=30., h=0.002):
    y=np.array([1.,0.]); ts=[]; xs=[]; resets=[]; last=-1e9; prev_e=y[0]
    for i in range(int(T/h)+1):
        t=i*h; ts.append(t); xs.append(y[0]); e=y[0]
        if reset and prev_e*e < 0 and t-last >= dwell:
            y[1]=0.; resets.append(t); last=t
        prev_e=e; y=rk4_step(y,h,lam,kp,ki)
    return np.asarray(ts),np.asarray(xs),resets

def discrete(lam, kp, ki, eta=.01, dwell=20, steps=3000, reset=True, noise=0.):
    th=1.; z=0.; old=None; last=-10**9; xs=[]; rs=[]
    rng=np.random.default_rng(123)
    for k in range(steps):
        g=lam*th + noise*rng.normal()
        z += g
        if reset and old is not None and old*g <= 0 and k-last >= dwell:
            z=0.; last=k; rs.append(k)
        th -= eta*(kp*g+ki*z)
        xs.append(th); old=g
    return np.asarray(xs),rs

def main():
    # Prediction 1: discriminant changes sign at R=kP^2*lambda/(4*kI)=1.
    lam=1.; ki=.4; Rvals=np.array([.25,.5,.75,1.,1.25,1.5,2.,3.])
    transition=[]
    for R in Rvals:
        kp=math.sqrt(4*ki*R/lam); rr=roots(lam,kp,ki)
        observed='oscillatory' if abs(rr[0].imag)>1e-10 else 'nonoscillatory'
        transition.append({'R':float(R),'predicted': 'oscillatory' if R<1 else 'nonoscillatory','observed':observed,'roots':[[float(x.real),float(x.imag)] for x in rr]})
    # Prediction 2: underdamped decay envelope rate is -kp*lambda/2.
    decay=[]
    for R in [.25,.5,2.,4.]:
        kp=math.sqrt(4*ki*R); rr=roots(lam,kp,ki); pred=-kp*lam/2
        t,x,_=continuous(lam,kp,ki,T=18,h=.001)
        # Fit log of local maxima magnitude, excluding initial transient and tiny values.
        peaks=[]
        for i in range(1,len(x)-1):
            if abs(x[i])>=abs(x[i-1]) and abs(x[i])>=abs(x[i+1]) and abs(x[i])>1e-7:
                peaks.append((t[i],abs(x[i])))
        if len(peaks)>=3:
            q=np.array(peaks[-min(8,len(peaks)):]); slope=np.polyfit(q[:,0],np.log(q[:,1]),1)[0]
        else: slope=float('nan')
        decay.append({'R':R,'predicted_rate':float(pred),'observed_rate':float(slope),'root_real_parts':[float(v.real) for v in rr]})
    # Prediction 3: resetting integral memory lowers post-crossing excursion; dwell suppresses chatter.
    kp=math.sqrt(4*ki*.25)
    no_t,no_x,no_r=continuous(lam,kp,ki,reset=False,T=20,h=.001)
    re_t,re_x,re_r=continuous(lam,kp,ki,reset=True,dwell=.0,T=20,h=.001)
    # peak absolute theta after first zero crossing, over next 1/4 period-ish 4 sec
    def post_peak(t,x):
        ix=np.flatnonzero(np.signbit(x[1:]) != np.signbit(x[:-1]))
        j=(ix[0]+1) if len(ix) else 0
        return float(np.max(np.abs(x[j:j+4000]))),float(t[j])
    p0,cross=post_peak(no_t,no_x); p1,_=post_peak(re_t,re_x)
    dwell_rows=[]
    for d in [0,20,50,100,200,500]:
        x,r=discrete(lam,kp,ki,eta=.01,dwell=d,steps=2500,reset=True,noise=.0)
        dwell_rows.append({'dwell_steps':d,'resets':len(r),'final_abs_theta':float(abs(x[-1])),'max_abs_theta':float(np.max(abs(x)))})
    # Small same-budget comparison on deterministic/stochastic scalar quadratic.
    comp={}
    xsg,_=discrete(lam,1.,0.,eta=.01,steps=2500,reset=False)
    xpi,_=discrete(lam,kp,ki,eta=.01,dwell=20,steps=2500,reset=False)
    xpir,rrr=discrete(lam,kp,ki,eta=.01,dwell=20,steps=2500,reset=True)
    comp={'SGD_final_loss':float(.5*xsg[-1]**2),'PI_final_loss':float(.5*xpi[-1]**2),'reset_PI_final_loss':float(.5*xpir[-1]**2),'reset_PI_resets':len(rrr),'SGD_steps_to_loss_1e-6':int(np.argmax(.5*xsg**2<1e-6)) if np.any(.5*xsg**2<1e-6) else None,'reset_PI_steps_to_loss_1e-6':int(np.argmax(.5*xpir**2<1e-6)) if np.any(.5*xpir**2<1e-6) else None}
    result={'claims':{'critical_boundary_R':1.0,'decay_rate_formula':'-kP*lambda/2 for underdamped modes','reset_effect':'lower post-crossing peak','dwell_effect':'fewer resets as dwell increases'},'critical_sweep':transition,'decay_sweep':decay,'reset_peak':{'no_reset_peak':p0,'reset_peak':p1,'first_crossing_time':cross,'reset_count':len(re_r)},'dwell_sweep':dwell_rows,'comparison':comp}
    OUT.write_text(json.dumps(result,indent=2))
    print(json.dumps(result,indent=2))
if __name__=='__main__': main()