import numpy as np
import random
import scipy.linalg as LA


N_ref=20
random.seed("uninitialised")

def init(seed):
    random.seed(seed)

def GOE(N=N_ref):
    mat=np.zeros((N,N))
    for i in xrange(N):
	mat[i,i]=random.normalvariate(0.0,1.0)
	for j in xrange(i+1,N):
	    mat[i,j]=random.normalvariate(0.0,1.0/np.sqrt(2))
	    mat[j,i]=mat[i,j]
    return mat

def Density_fun(N=N_ref):
    def fun(x):
	if abs(x)>=np.sqrt(2*N):
	    return 0
	else:
	    return (np.sqrt(2*N-x**2))/np.pi
    return fun

def Cum_Density_fun(N=N_ref):
    def fun(x):
	if x<=-np.sqrt(2*N):
	    return float(-N/2.0)
	elif x>=np.sqrt(2*N):
	    return float(N/2.0)
	else:
	    fac=0.5*x*np.sqrt(2*N-x**2)+N*np.arctan(x/np.sqrt(2*N-x**2))
	    return fac/np.pi
    return fun

def GOE_eigs(N=N_ref,mat=None):
    if not mat:
	mat=GOE(N=N)
    eig=LA.eigh(mat,eigvals_only=True)
    return eig

def GOE_unfolded(N=N_ref,eigs=None,mat=None,cum_func=None):
    if not eigs:
	eigs=GOE_eigs(N=N,mat=mat)
    if not cum_func:
	cum_func=Cum_Density_fun(N=N)
    return unfold(eigs,cum_func)

def unfold(states,cum_func):
    unfolded=np.zeros_like(states)
    for i in xrange(len(states)):
	unfolded[i]=cum_func(states[i])
    return unfolded

def NNS(unfolded,calc_range=(-np.inf,np.inf)):
    NNS=[]
    for i in range(len(unfolded)-1):
	if unfolded[i]<calc_range[0]:
	    continue
	elif unfolded[i+1]>calc_range[1]:
	    continue
        else:
	    NNS.append(unfolded[i+1]-unfolded[i])
    return NNS

#NN=20
#for i in xrange(5):
#    vals=GOE_unfolded(NN)
#    spacings=NNS(vals,(-0.5*NN/2.0,0.5*NN/2.0))
#    print np.average(spacings)









