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

SEED = 2313

def set_seed(seed=SEED):
    random.seed(seed); np.random.seed(seed)

def metrics(x, dt=1.0):
    x=np.asarray(x,float); z=x-x.mean(); v=float(np.mean(z*z))
    rho=float(np.sum(z[1:]*z[:-1])/max(np.sum(z*z),1e-15))
    return v,rho,float(-math.log(max(rho,1e-3))/dt)

def ar1_sweep():
    rng=np.random.default_rng(SEED); a_values=np.array([.2,.4,.6,.75,.85,.92,.96,.98])
    sigma=.5; n=120000; burn=1000; rows=[]
    for a in a_values:
        e=rng.normal(0,sigma,n+burn); x=np.zeros(n+burn)
        for i in range(1,len(x)): x[i]=a*x[i-1]+e[i]
        v,rho,r=metrics(x[burn:]); rows.append({'a':float(a),'rho_hat':rho,'rho_pred':float(a),
          'var_hat':v,'var_pred':float(sigma*sigma/(1-a*a)),'recovery_hat':r,
          'recovery_pred':float(-math.log(a))})
    return {'rows':rows,'rho_rmse':float(np.sqrt(np.mean([(q['rho_hat']-q['rho_pred'])**2 for q in rows]))),
      'variance_mean_relative_error':float(np.mean([abs(q['var_hat']-q['var_pred'])/q['var_pred'] for q in rows])),
      'recovery_rmse':float(np.sqrt(np.mean([(q['recovery_hat']-q['recovery_pred'])**2 for q in rows])))}

def delayed_quadratic(scheduler, seed=SEED, steps=3000, delay=8, lr=.035, momentum=.85):
    A=np.diag([1.,10.8]); x=np.array([2.,2.]); velocity=np.zeros(2); hist=[]
    lr_now=lr; beta_now=momentum; triggered=False; trigger_step=None; ema=None; vals=[]
    base_stats=None; cooldown=0; losses=[]
    for t in range(steps):
        loss=.5*float(x@A@x); g=A@x
        stale=g if delay==0 or t<delay else hist[t-delay]
        velocity=beta_now*velocity+stale; x=x-lr_now*velocity; hist.append(g.copy()); losses.append(loss)
        # Fixed random parameter projection is an alternative scalar observable.
        obs=float(np.log1p(loss)); ema=obs if ema is None else .97*ema+.03*obs; vals.append(obs-ema)
        if len(vals)>=60 and t%10==0:
            v,rho,rate=metrics(vals[-60:])
            if base_stats is None and t==59: base_stats=(max(v,1e-10),rho,rate)
            if scheduler and not triggered and base_stats is not None and t>=70:
                bv,br,bc=base_stats
                if rho>br+.10 or v>2*bv or rate<.65*bc:
                    lr_now*=.5; beta_now=min(beta_now,.90); triggered=True; trigger_step=t
            if triggered and cooldown==0 and rate>.9*base_stats[2]:
                lr_now=lr; beta_now=momentum; cooldown=100
        cooldown=max(0,cooldown-1)
        if not np.all(np.isfinite(x)) or np.linalg.norm(x)>1e8:
            return {'diverged':True,'divergence_step':t,'final_loss':'inf','trigger_step':trigger_step,'min_loss':float(np.min(losses))}
    return {'diverged':False,'divergence_step':None,'final_loss':float(losses[-1]),'trigger_step':trigger_step,'min_loss':float(np.min(losses))}

def training_sweep():
    rows=[]
    for d in [0,2,4,6,8,10,12]:
        rows.append({'delay':d,'baseline':delayed_quadratic(False,delay=d),'scheduler':delayed_quadratic(True,delay=d)})
    return rows

def main():
    set_seed(); out={'seed':SEED,'ar1':ar1_sweep(),'delayed_training':training_sweep()}
    Path('results.json').write_text(json.dumps(out,indent=2)); print(json.dumps(out,indent=2))
if __name__=='__main__': main()