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

SEED = 2845
NU = 5.0
EPS = 2e-3

# Student-t target: U(x)=(nu+1)/2 log(1+x^2/nu), so Var[X]=nu/(nu-2).
def U(x):
    return 0.5*(NU+1.0)*np.log1p(x*x/NU)

def grad_U(x):
    return (NU+1.0)*x/(NU+x*x)

def sigma(x, alpha):
    r = np.sqrt(x*x + 1e-8)
    return 1.0 + alpha*np.log1p(r)

def dsigma2(x, alpha):
    r = np.sqrt(x*x + 1e-8)
    s = 1.0 + alpha*np.log1p(r)
    return 2.0*s*alpha*x/(r*(1.0+r))

def drift(x, alpha, correction=True):
    s = sigma(x, alpha)
    return (dsigma2(x, alpha) if correction else 0.0) - s*s*grad_U(x)

def current_residual(x, alpha):
    # J=b*pi-d(a*pi)/dx; analytic derivative gives d(a*pi)=(a'-a U')pi.
    pi=np.exp(-U(x)); a=sigma(x,alpha)**2
    return drift(x,alpha,True)*pi - (dsigma2(x,alpha)-a*grad_U(x))*pi

def run(alpha, correction=True, nsteps=50000, nchains=128, burn=8000, seed=0):
    rng=np.random.default_rng(seed)
    x=rng.normal(0,2,nchains)
    samples=[]; grad_calls=0
    for k in range(nsteps):
        s=sigma(x,alpha)
        x=x+EPS*drift(x,alpha,correction)+np.sqrt(2*EPS)*s*rng.normal(size=nchains)
        grad_calls += nchains
        if k>=burn and k%5==0: samples.append(x.copy())
    z=np.concatenate(samples)
    # Batch means gives a conservative ESS estimate.
    flat=z.reshape(-1); m=100
    nb=len(flat)//m
    bm=flat[:nb*m].reshape(nb,m).mean(1)
    var=np.var(flat); ess=min(len(flat), nb*m*var/(m*np.var(bm)+1e-30))
    return {'mean':float(np.mean(flat)), 'second':float(np.mean(flat**2)),
            'q95':float(np.quantile(flat,.95)), 'ess':float(ess),
            'ess_per_grad':float(ess/grad_calls), 'max_abs':float(np.max(np.abs(flat)))}

def main():
    # Prediction 1: divergence correction makes stationary current identically zero.
    xs=np.linspace(-12,12,10001)
    current={str(a):float(np.max(np.abs(current_residual(xs,a)))) for a in [0,.25,.5,1.0,2.0]}

    # Prediction 2: sigma(1,alpha)-1 is exactly linear in alpha.
    alphas=np.array([0,.25,.5,1.,2.])
    sig1=np.array([sigma(np.array(1.),a) for a in alphas])
    slope=float(np.polyfit(alphas,sig1,1)[0])
    linear_maxerr=float(np.max(np.abs(sig1-(1+alphas*np.log(2)))))

    # Prediction 3: correction magnitude is alpha*(1+alpha*c), hence quadratic
    # coefficient is predicted from analytic expansion at x=2.
    x0=2.; h=1e-4
    vals=[]
    for a in alphas:
        vals.append(abs(dsigma2(np.array(x0),a)))
    vals=np.array(vals)
    fit=np.polyfit(alphas,vals,2)
    c=np.log1p(abs(x0)); predicted_quad=2*c*(x0/(abs(x0)*(1+abs(x0))))
    # Direct finite-difference confirmation of the derivative correction.
    fd=[]
    for a in [.25,.5,1.0]:
        f=lambda q:sigma(np.array(q),a)**2
        fd.append(abs((f(x0+h)-f(x0-h))/(2*h)-dsigma2(np.array(x0),a)))
    fd_max=float(max(fd))

    # Mechanism sweep and baseline: same ULA update, with/without correction.
    rows=[]
    for a in [0.,.5,1.0,2.0]:
        idea=run(a,True,seed=SEED+int(100*a))
        base=run(a,False,seed=SEED+int(100*a))
        rows.append({'alpha':a,'corrected':idea,'uncorrected':base})

    out={'target_variance':NU/(NU-2), 'predictions':{
        'zero_current_max_abs':current,
        'linear_sigma_at_x1':{'predicted_slope':math.log(2),'observed_slope':slope,'max_error':linear_maxerr},
        'quadratic_correction_at_x2':{'fit_coefficients_constant_linear_quadratic':fit.tolist(), 'predicted_quadratic_coefficient':predicted_quad, 'fd_max_error':fd_max}},
        'sweep':rows}
    Path('results.json').write_text(json.dumps(out,indent=2))
    print(json.dumps(out,indent=2))

if __name__=='__main__': main()