import numpy as np


def energy(z):
    n = len(z)
    d = z[:, None, :] - z[None, :, :]
    r = np.sqrt(np.sum(d*d, axis=-1))
    iu = np.triu_indices(n, 1)
    return 0.5*np.sum(z*z) - np.sum(np.log(r[iu]))/n


def force(z, eps=1e-5):
    # negative gradient of H: -z + repulsive Coulomb term
    d = z[:, None, :] - z[None, :, :]
    r2 = np.sum(d*d, axis=-1) + eps
    np.fill_diagonal(r2, np.inf)
    return -z + np.sum(d/r2[:, :, None], axis=1)/len(z)


def finite_difference_check(seed=0):
    rng=np.random.default_rng(seed)
    z=rng.normal(size=(7,2)); v=rng.normal(size=z.shape); v/=np.linalg.norm(v)
    vals=[]
    for h in [1e-2, 1e-3, 3e-4, 1e-4]:
        fd=(energy(z+h*v)-energy(z-h*v))/(2*h)
        analytic=-np.sum(force(z,eps=0.0)*v)
        vals.append((h, abs(fd-analytic), abs(fd), abs(analytic)))
    return vals


def langevin(n, steps, beta=2.0, eta=.015, seed=0, eps=1e-5, burn=0, thin=1):
    rng=np.random.default_rng(seed)
    a=np.linspace(0,2*np.pi,n,endpoint=False)+rng.normal(0,.03,n)
    z=np.c_[np.cos(a),np.sin(a)]*np.sqrt(n/2.0)
    out=[]
    for t in range(steps):
        z += eta*force(z,eps) + np.sqrt(2*eta/beta)*rng.normal(size=z.shape)
        if t>=burn and (t-burn)%thin==0: out.append(z.copy())
    return np.asarray(out)


def acf_iat(x, maxlag=None):
    x=np.asarray(x,float); x=x-x.mean(); n=len(x)
    if n<10 or np.var(x)<1e-14:return np.nan
    maxlag=min(n//3, maxlag or n//3)
    lags=np.arange(maxlag+1)
    # Exactly maxlag+1 entries: lags 0,...,maxlag.
    ac=np.correlate(x,x,mode='full')[n-1:n+maxlag]/(n-lags)
    ac/=ac[0]
    k=1
    while k<len(ac) and ac[k]>0:k+=1
    return float(max(1.,1+2*np.sum(ac[1:k])))


def bank_stats(z, seed=123):
    rng=np.random.default_rng(seed)
    d=np.sqrt(np.sum((z[:,None]-z[None,:])**2,axis=-1)); np.fill_diagonal(d,np.inf)
    nn=d.min(axis=1)
    q=rng.normal(size=(4000,2))*np.sqrt(2)
    assign=np.argmin(np.sum((q[:,None,:]-z[None,:,:])**2,axis=-1),axis=1)
    counts=np.bincount(assign,minlength=len(z))/len(q)
    entropy=-np.sum(counts[counts>0]*np.log(counts[counts>0]))/np.log(len(z))
    dead=float(np.mean(counts==0))
    return float(np.median(nn)), float(np.min(nn)), float(entropy), dead


def main():
    print('FINITE_DIFFERENCE h abs_error abs_fd abs_analytic')
    for row in finite_difference_check():print('%g %.3e %.3e %.3e'%row)
    print('MIXING n iat_radius iat_mean_radius median_nn min_nn entropy dead')
    for n in [8,32,128]:
        traj=langevin(n, steps=4500, burn=1000, thin=1, seed=10+n)
        radius=np.linalg.norm(traj,axis=2)
        iat=np.nanmean([acf_iat(radius[:,i]) for i in range(n)])
        mean_iat=acf_iat(radius.mean(axis=1))
        st=bank_stats(traj[-1],seed=90+n)
        print(n, '%.2f %.2f %.4f %.4f %.4f %.4f'%(iat,mean_iat,*st))
    print('GAUSSIAN_CONTROL n median_nn min_nn entropy dead')
    for n in [8,32,128]:
        rng=np.random.default_rng(100+n); z=rng.normal(size=(n,2))*np.sqrt(n/2)
        print(n,'%.4f %.4f %.4f %.4f'%bank_stats(z,seed=90+n))

if __name__=='__main__': main()