import json, math, random
import numpy as np
import mpmath as mp

SEED=1234
np.random.seed(SEED); random.seed(SEED)

def kappa(m,beta):
    if m*beta <= 1: raise ValueError('requires m beta > 1')
    return beta*beta/(8*(m*beta-1)*(2*m+1))

def V(a):
    a=np.asarray(a); return sum((a[i]-a[j])**2 for i in range(len(a)) for j in range(i+1,len(a)))

def rho_det_mp(x):
    mp.mp.dps=80
    x=[mp.mpf(str(v)) for v in x]; m=len(x)
    A=mp.matrix(m)
    for i in range(m):
        for j in range(m):
            d=x[i]-x[j]
            A[i,j]=1/(2*mp.pi) if d==0 else mp.sin(d/2)/(mp.pi*d)
    return mp.det(A)

def normalized_ratio_mp(e,a,beta=2):
    e=mp.mpf(str(e)); aa=[mp.mpf(str(v)) for v in a]; m=len(aa)
    lead=abs(e)**(beta*m*(m-1)/2)
    for i in range(m):
        for j in range(i+1,m): lead*=abs(aa[i]-aa[j])**beta
    return rho_det_mp([e*v for v in aa])/lead

def prior_score(x,beta=2,delta=.02,include_correction=True):
    x=np.asarray(x,float); m=len(x); kap=kappa(m,beta) if include_correction else 0.
    d=x[:,None]-x[None,:]; den=d*d+delta*delta
    rep=np.divide(d,den,out=np.zeros_like(d),where=den>0)
    np.fill_diagonal(rep,0); np.fill_diagonal(d,0)
    return beta*np.sum(rep,axis=1)-2*kap*np.sum(d,axis=1)

def langevin(kind, n=8, beta=2, steps=2800, burn=800, dt=.0015, delta=.03):
    x=np.sort(np.random.randn(n)); samples=[]
    for it in range(steps):
        if kind=='baseline': g=-x
        elif kind=='repulsion': g=-x+prior_score(x,beta,delta,False)
        else: g=-x+prior_score(x,beta,delta,True)
        x=x+dt*g+math.sqrt(2*dt)*np.random.randn(n)
        if it>=burn: samples.append(np.sort(x.copy()))
    return np.asarray(samples)

def metrics(samples):
    gaps=np.diff(samples,axis=1).ravel(); gaps=gaps/(np.mean(gaps)+1e-12)
    return {'collision_rate_gap_lt_0.10':float(np.mean(gaps<.10)), 'gap_std':float(np.std(gaps)), 'min_gap_p05':float(np.quantile(gaps,.05))}

def theorem_check(a):
    m=len(a); kap=kappa(m,2); es=np.array([.004,.006,.009,.013,.018,.025,.035,.05])
    rs=np.array([float(normalized_ratio_mp(e,a)) for e in es])
    # R(e)=C + B e^2 + D e^4; theorem predicts -B/(C V)=kappa.
    z=es**2; A=np.column_stack([np.ones(len(z)),z,z*z])
    C,B,_=np.linalg.lstsq(A,rs,rcond=None)[0]
    est=-B/(C*V(a))
    return {'m':m,'theory_kappa':kap,'estimated_kappa':float(est),'relative_error':float(abs(est-kap)/kap), 'leading_constant':float(C)}

def main():
    checks=[theorem_check([0.,1.]),theorem_check([0.,1.,2.])]
    x=np.array([-.71,-.11,.63]); b=1.7; h=1e-6
    def logp(y): return b*sum(math.log(abs(y[i]-y[j])) for i in range(3) for j in range(i+1,3))-kappa(3,b)*V(y)
    fd=np.array([(logp(x+np.eye(3)[i]*h)-logp(x-np.eye(3)[i]*h))/(2*h) for i in range(3)])
    analytic=prior_score(x,b,0.0,True)
    score={'max_abs_error':float(np.max(abs(fd-analytic))),'finite_difference':fd.tolist(),'analytic':analytic.tolist()}
    sam={k:metrics(langevin(k)) for k in ['baseline','repulsion','full']}
    result={'seed':SEED,'theorem_checks':checks,'score_check':score,'sampler':sam}
    print(json.dumps(result,indent=2)); open('results.json','w').write(json.dumps(result,indent=2))

if __name__=='__main__': main()