import math
import numpy as np

SEED = 2517
A, B, KAPPA = 0.8, 3.0, 0.35
XMIN = 0.20

def phi(x): return np.sin(2*np.pi*x)
def phip(x): return 2*np.pi*np.cos(2*np.pi*x)
def W(x, n):
    x=np.asarray(x,float); y=np.zeros_like(x)
    for j in range(n+1): y += A**j*phi(B**j*x)
    return y
def Wprime(x,n):
    x=np.asarray(x,float); y=np.zeros_like(x)
    for j in range(n+1): y += (A*B)**j*phip(B**j*x)
    return y
def dq_direct(x,n): return (W(x,n)-A*W(B*x,n))/x
def dq_closed(x,n): return (phi(x)-A**(n+1)*phi(B**(n+1)*x))/x
def V(x): return .5*(x-1.5)**2
def Vp(x): return x-1.5
def U(x,n): return V(x)+KAPPA*W(x,n)
def force(x,n,kind):
    return Vp(x) + KAPPA*(Wprime(x,n) if kind=='classical' else dq_closed(x,n))

def proposal(x,p,n,kind,eps,L):
    x0,p0=x,p
    for _ in range(L):
        p -= .5*eps*force(x,n,kind); x += eps*p
        if x <= XMIN or not np.isfinite(x): return x0,p0,False
        p -= .5*eps*force(x,n,kind)
    dH=U(x,n)+.5*p*p-U(x0,n)-.5*p0*p0
    ok=math.log(np.random.random()) < min(0.,-dH)
    return (x,p,True) if ok else (x0,p0,False)

def chain(n,kind,eps=.01,L=4,steps=5000,burn=1000):
    x=1.5; xs=[]; accepted=0; valid=0
    for t in range(steps):
        p=np.random.normal(); xn,pn,ok=proposal(x,p,n,kind,eps,L)
        valid += int(ok)
        if ok: accepted += int(xn != x or pn != p); x=xn
        if t>=burn: xs.append(x)
    z=np.asarray(xs); zc=z-z.mean(); ac1=np.dot(zc[:-1],zc[1:])/max(1e-12,np.dot(zc,zc))
    ess=len(z)*(1-ac1)/(1+ac1)
    return float(accepted/steps),float(valid/steps),float(z.mean()),float(ess)

def main():
    np.random.seed(SEED)
    grid=np.linspace(XMIN,4.0,4000)
    print('CORE_MATH')
    for n in [0,2,4,8,12,16]:
        err=float(np.max(np.abs(dq_direct(grid,n)-dq_closed(grid,n))))
        md=float(np.max(np.abs(dq_closed(grid,n))))
        mp=float(np.max(np.abs(Wprime(grid,n))))
        print(f'N={n:2d} closure_maxerr={err:.3e} quotient_max={md:.3f} derivative_max={mp:.3f}')
    print('SCALING_PREDICTION')
    # Predicted quotient envelope <= 2/xmin; derivative leading envelope grows (ab)^N.
    for n in [2,4,6,8,10,12,14,16]:
        q=float(np.percentile(np.abs(dq_closed(grid,n)),99))
        d=float(np.percentile(np.abs(Wprime(grid,n)),99))
        print(f'N={n:2d} q99={q:.3f} classical_dq99={d:.3f} ratio={d/max(q,1e-9):.2f}')
    print('SAMPLER eps=.01 L=4, exact target MH correction')
    for n in [4,8,12,16]:
        np.random.seed(SEED+n)
        c=chain(n,'classical'); np.random.seed(SEED+n)
        q=chain(n,'quotient')
        print(f'N={n:2d} classical acc={c[0]:.3f} mean={c[2]:.3f} ESS={c[3]:.1f} quotient acc={q[0]:.3f} mean={q[2]:.3f} ESS={q[3]:.1f}')
    print('STEP_SWEEP_N16')
    for eps in [.002,.005,.01,.02]:
        np.random.seed(SEED+int(eps*10000)); c=chain(16,'classical',eps=eps,steps=2500,burn=500)
        np.random.seed(SEED+int(eps*10000)); q=chain(16,'quotient',eps=eps,steps=2500,burn=500)
        print(f'eps={eps:.3f} classical_acc={c[0]:.3f} quotient_acc={q[0]:.3f}')

if __name__=='__main__': main()