import json
import numpy as np

A, Q, R, SIGMA2 = 1.0, 0.80, 1.0, 0.25
SIGMA = np.sqrt(SIGMA2)

def grad(x): return A*x + Q*x*x/2.0 + R*x*x*x

def theoretical_bias(gamma, n, h):
    hh=(h-1)*(5*h-1)/h
    return -Q*SIGMA2/(4*A*A)*gamma/n - Q*SIGMA2/(12*A)*hh*gamma*gamma

class BiasEstimator:
    def __init__(self, delta=0.02, ema=0.08, clip=0.5):
        self.delta,self.ema,self.clip=delta,ema,clip
        self.f2,self.f3,self.var=A,0.0,SIGMA2
    def update(self,x,residuals,gamma,n,h):
        d=self.delta
        f2n=(grad(x+d)-grad(x-d))/(2*d)
        # Since grad=f', f''' is the second derivative of grad.
        f3n=(grad(x+d)-2*grad(x)+grad(x-d))/(d*d)
        varn=float(np.var(residuals)); r=self.ema
        self.f2=(1-r)*self.f2+r*f2n; self.f3=(1-r)*self.f3+r*f3n; self.var=(1-r)*self.var+r*varn
        aa=max(self.f2,.05); hh=(h-1)*(5*h-1)/h
        b=-self.f3*self.var/(4*aa*aa)*gamma/n-self.f3*self.var/(12*aa)*hh*gamma*gamma
        return float(np.clip(b,-self.clip,self.clip))

def simulate(n,h,gamma,mode,seed,rounds=8000,burn=2000):
    rng=np.random.default_rng(seed); x=0.; xi=np.zeros(n); est=BiasEstimator(); next_c=0.; xs=[]; reported=[]; losses=[]
    for t in range(rounds):
        theta=np.full(n,x); residuals=[]
        for _ in range(h):
            eps=rng.normal(0,SIGMA,n); residuals.append(eps)
            theta-=gamma*(grad(theta)+xi+ (next_c if mode=='gradient_correction' and t>=1000 else 0.) + eps)
        x=float(theta.mean()); xi += (theta-x)/(gamma*h)
        b=est.update(x,np.concatenate(residuals),gamma,n,h)
        if mode=='gradient_correction' and t>=1000: next_c=est.f2*b
        if t>=burn:
            xs.append(x); reported.append(x-b if mode=='output_correction' else x)
            losses.append(A*x*x/2+Q*x*x*x/6)
    xs=np.asarray(xs); reported=np.asarray(reported)
    return {'mean_x':float(xs.mean()),'abs_bias':float(abs(xs.mean())),'std_x':float(xs.std()),
            'mean_loss':float(np.mean(losses)),'mean_reported_x':float(reported.mean()),
            'theory':float(theoretical_bias(gamma,n,h))}

def math_check():
    est=BiasEstimator(delta=.01,ema=1.)
    b=est.update(0.,np.random.default_rng(1).normal(0,SIGMA,100000),.01,64,8)
    return {'f2_hat':est.f2,'f3_hat':est.f3,'sigma2_hat':est.var,'expected_f2':A,'expected_f3':Q,
            'formula_b':b,'formula_exact':theoretical_bias(.01,64,8)}

def main():
    out={'math_check':math_check(),'experiments':{}}
    for h in (8,):
      for n in (8,64):
        k=f'n{n}_h{h}'; seed=100+n+h
        out['experiments'][k]={'vanilla':simulate(n,h,.01,'vanilla',seed),
          'output_correction':simulate(n,h,.01,'output_correction',seed),
          'gradient_correction':simulate(n,h,.01,'gradient_correction',seed)}
    out['scaling_large_N']=[]
    for gamma in (.012,.020,.028):
      r=simulate(64,8,gamma,'vanilla',777+int(gamma*10000),rounds=8000,burn=2000)
      out['scaling_large_N'].append({'gamma':gamma,'empirical_bias':r['mean_x'],'theory':r['theory']})
    with open('results.json','w') as f: json.dump(out,f,indent=2)
    print(json.dumps(out,indent=2))
if __name__=='__main__': main()