import sys
import math
from err_function import *

#from Alex Travesset trvsst@ameslab.gov
#last updated February 2010

def Ewald_Oblique(x1,y1,z1,q1,alpha,b_1,c_1,r_cut):
#computes the Ewald sum for a 3D system with the 2D periodicity of
#an oblique lattice, the z-coordinate is NOT periodic
#the unit vectors are
# a_1=(a,0) a_2=(c*a,b*a)
#The area is A=b*a^2
#square lattice is c=0 and b=1
#triangular lattice is c=1/2 b=sqrt(3)/2

#x1,y1,z1 and q1 are a list containing the x,y,z and charge of the particles
#alpha is the split parameter
#note that the lattice constant is taken as 1
#and the oblique lattice is defined by unit vectors (1,0) (c_1,b_1)
#r_cut is the cut-off value

#returns a list of 4 elements storing the electrostatic energy contribution
#   Ewald_c[0]=Total contribution
#   Ewald_c[1]=Real space contribution
#   Ewald_c[2]=Reciprocal space contribution
#   Ewald_c[3]=Additional contributions along the z axis and constant sums

	Ewald_c=[0,0,0,0]

#make sure that c_1 is positive

	if(c_1<0):
		c_1=-c_1
#The area of the lattice unit cell is
	
	Ac=b_1
		
#First make sure that all the particles specified are within the unit cell, that is

	eps=1e-10
	for i in range(len(x1)):
		t1=x1[i]-y1[i]*c_1/b_1
		t2=y1[i]/b_1
		if((t1>1+eps)or(t1<-eps)or(t2>1+eps)or(t2<-eps)):
			print 'x,y coordinates not within the primitive unit cell'
			print t1,t2
			print x1[i],y1[i]
			sys.exit()

	Q=sum(q1)

#Output a warning if the system is not electrically neutral
	
	if(Q!=0):
		print 'Non-zero net charge, added a constant background to keep system neutral'
	
#set precision
	eps=1e-11;
#Now compute the real space contribution
#the maximum value of the integers defining the Ewald sum

	n1_max=1+int(math.ceil((1+c_1/b_1)*r_cut/alpha))
	n2_max=1+int(math.ceil(r_cut/(alpha*b_1)))

	En_real=0.0;

	for i1 in range(len(x1)):
		for i2 in range(i1,len(x1),1):
			fac=1.0
			if(i1==i2):
				fac=0.5
			for j1 in range(-n1_max,n1_max+1,1):
				for j2 in range(-n2_max,n2_max+1,1):	
					del_x=x1[i1]-x1[i2]+j1+c_1*j2
					del_y=y1[i1]-y1[i2]+b_1*j2
					del_z=z1[i1]-z1[i2]
					arg_r=alpha*math.sqrt(del_x*del_x+del_y*del_y+del_z*del_z)
					if( (arg_r>eps)and(arg_r<r_cut)	):
						En_real+=fac*q1[i1]*q1[i2]*alpha*(1-erf(arg_r))/arg_r


	Ewald_c[0]=En_real
	Ewald_c[1]=En_real

#Compute the reciprocal space contribution

	En_recp=0.0
	M_PI=math.pi

	for i1 in range(len(x1)):
		for i2 in range(i1,len(x1),1):
			fac=1.0
			if(i1==i2):
				fac=0.5

			del_x=x1[i1]-x1[i2]
			del_y=y1[i1]-y1[i2]
			h_del_z=z1[i1]-z1[i2]
		
			del_z=abs(h_del_z)

			n1_max=int(math.ceil((r_cut-del_z)*2*alpha/(M_PI*math.sqrt(1+c_1*c_1/(b_1*b_1)))))
			n2_max=int(math.ceil((r_cut-del_z)*2*alpha*b_1/M_PI))

			if(n1_max<0):
				n1_max=0

			if(n2_max<0):
				n2_max=0

			for j1 in range(-n1_max,n1_max+1,1):
				for j2 in range(-n2_max,n2_max+1,1):	
					G_x=2*M_PI*j1
					G_y=2*M_PI*(j2-c_1*j1)/b_1
					mod_G=math.sqrt(G_x*G_x+G_y*G_y)
					arg_G=mod_G/(2*alpha)+alpha*del_z
					r_G=G_x*del_x+G_y*del_y
					if( (mod_G>eps)and(arg_G<r_cut)	):
						En_recp+=fac*q1[i1]*q1[i2]*(1-erf(arg_G))*math.exp(del_z*mod_G)*math.cos(r_G)/mod_G

			del_z=-del_z

			n1_max=int(math.ceil((r_cut-del_z)*2*alpha/(M_PI*math.sqrt(1+c_1*c_1/(b_1*b_1)))))
			n2_max=int(math.ceil((r_cut-del_z)*2*alpha*b_1/M_PI))

			if(abs(del_z)>abs(math.log((1-erf(r_cut))))):
				n1_max=0
				n2_max=0
		
			for j1 in range(-n1_max,n1_max+1,1):
				for j2 in range(-n2_max,n2_max+1,1):	
					G_x=2*M_PI*j1
					G_y=2*M_PI*(j2-c_1*j1)/b_1
					mod_G=math.sqrt(G_x*G_x+G_y*G_y)
					arg_G=mod_G/(2*alpha)+alpha*del_z
					r_G=G_x*del_x+G_y*del_y
					if( (mod_G>eps)and(arg_G<r_cut)	):
						En_recp+=fac*q1[i1]*q1[i2]*(1-erf(arg_G))*math.exp(del_z*mod_G)*math.cos(r_G)/mod_G


	Ewald_c[0]+=M_PI*En_recp/Ac
	Ewald_c[2]=M_PI*En_recp/Ac

#Now compute the remaining terms to complete the Ewald summation

	En_rest=0.0

	for i1 in range(len(x1)):
		for i2 in range(len(x1)):
			del_z=alpha*(z1[i1]-z1[i2])
			En_rest+=q1[i1]*q1[i2]*(math.exp(-del_z*del_z)+math.sqrt(M_PI)*del_z*erf(del_z))/alpha
				
	En_rest*=math.sqrt(M_PI)/Ac

	En_charge=0;
	for i in range(len(x1)):
		En_charge+=q1[i]*q1[i]

	En_rest+=alpha*En_charge/math.sqrt(M_PI)

	Ewald_c[0]+=-En_rest
	Ewald_c[3]=-En_rest
	
	return Ewald_c


def Ewald_Oblique_Image(x1,y1,z1,q1,alpha,b_1,c_1,r_cut):
#computes the Ewald sum for the image charges of a 3D system with 
#the 2D periodicity of
#an oblique lattice, the z-coordinate is NOT periodic
#It is assumed that the boundary between two dielectric constants is
#at z=0, so all particles have coordinates z<0

#the unit vectors are
# a_1=(a,0) a_2=(c*a,b*a)
#The area is A=b*a^2
#square lattice is c=0 and b=1
#triangular lattice is c=1/2 b=sqrt(3)/2

#x1,y1,z1 and q1 are a list containing the x,y,z and charge of the particles
#alpha is the split parameter
#note that the lattice constant is taken as 1
#and the oblique lattice is defined by unit vectors (1,0) (c_1,b_1)
#r_cut is the cut-off value

#returns a list of 4 elements storing the electrostatic energy contribution
#   Ewald_c[0]=Total contribution
#   Ewald_c[1]=Real space contribution
#   Ewald_c[2]=Reciprocal space contribution
#   Ewald_c[3]=Additional contributions along the z axis and constant sums

	Ewald_c=[0,0,0,0]

#make sure that c_1 is positive

	if(c_1<0):
		c_1=-c_1
#The area of the lattice unit cell is
	
	Ac=b_1
		
#First make sure that all the particles specified are within the unit cell, that is

	for i in range(len(x1)):
		t1=x1[i]-y1[i]*c_1/b_1
		t2=y1[i]/b_1
		if((t1>1)or(t1<0)or(t2>1)or(t2<0)):
			print 'x,y coordinates not within the primitive unit cell'
			sys.exit()


#Make sure that all the particles are located at z<0

	for i in range(len(z1)):
		if(z1[i]>0):
			print 'particles must be located at z<0'
			sys.exit()


	Q=sum(q1)

#Output a warning if the system is not electrically neutral
	
	if(Q!=0):
		print 'Non-zero net charge, added a constant background to keep system neutral'
	
#set precision
	eps=1e-11;
#Now compute the real space contribution
#the maximum value of the integers defining the Ewald sum

	n1_max=1+int(math.ceil((1+c_1/b_1)*r_cut/alpha))
	n2_max=1+int(math.ceil(r_cut/(alpha*b_1)))

	En_real=0.0;

	for i1 in range(len(x1)):
		for i2 in range(i1,len(x1),1):
			fac=1.0
			if(i1==i2):
				fac=0.5
			for j1 in range(-n1_max,n1_max+1,1):
				for j2 in range(-n2_max,n2_max+1,1):	
					del_x=x1[i1]-x1[i2]+j1+c_1*j2
					del_y=y1[i1]-y1[i2]+b_1*j2
					del_z=z1[i1]+z1[i2]
					arg_r=alpha*math.sqrt(del_x*del_x+del_y*del_y+del_z*del_z)
					if(arg_r<r_cut):
						En_real+=fac*q1[i1]*q1[i2]*alpha*(1-erf(arg_r))/arg_r


	Ewald_c[0]=En_real
	Ewald_c[1]=En_real

#Compute the reciprocal space contribution

	En_recp=0.0
	M_PI=math.pi

	for i1 in range(len(x1)):
		for i2 in range(i1,len(x1),1):
			fac=1.0
			if(i1==i2):
				fac=0.5

			del_x=x1[i1]-x1[i2]
			del_y=y1[i1]-y1[i2]
			h_del_z=z1[i1]+z1[i2]
		
			del_z=abs(h_del_z)

			n1_max=int(math.ceil((r_cut-del_z)*2*alpha/(M_PI*math.sqrt(1+c_1*c_1/(b_1*b_1)))))
			n2_max=int(math.ceil((r_cut-del_z)*2*alpha*b_1/M_PI))

			if(n1_max<0):
				n1_max=0

			if(n2_max<0):
				n2_max=0

			for j1 in range(-n1_max,n1_max+1,1):
				for j2 in range(-n2_max,n2_max+1,1):	
					G_x=2*M_PI*j1
					G_y=2*M_PI*(j2-c_1*j1)/b_1
					mod_G=math.sqrt(G_x*G_x+G_y*G_y)
					arg_G=mod_G/(2*alpha)+alpha*del_z
					r_G=G_x*del_x+G_y*del_y
					if( (mod_G>eps)and(arg_G<r_cut)	):
						En_recp+=fac*q1[i1]*q1[i2]*(1-erf(arg_G))*math.exp(del_z*mod_G)*math.cos(r_G)/mod_G

			del_z=-del_z

			n1_max=int(math.ceil((r_cut-del_z)*2*alpha/(M_PI*math.sqrt(1+c_1*c_1/(b_1*b_1)))))
			n2_max=int(math.ceil((r_cut-del_z)*2*alpha*b_1/M_PI))

			if(abs(del_z)>abs(math.log((1-erf(r_cut))))):
				n1_max=0
				n2_max=0
		
			for j1 in range(-n1_max,n1_max+1,1):
				for j2 in range(-n2_max,n2_max+1,1):	
					G_x=2*M_PI*j1
					G_y=2*M_PI*(j2-c_1*j1)/b_1
					mod_G=math.sqrt(G_x*G_x+G_y*G_y)
					arg_G=mod_G/(2*alpha)+alpha*del_z
					r_G=G_x*del_x+G_y*del_y
					if( (mod_G>eps)and(arg_G<r_cut)	):
						En_recp+=fac*q1[i1]*q1[i2]*(1-erf(arg_G))*math.exp(del_z*mod_G)*math.cos(r_G)/mod_G


	Ewald_c[0]+=M_PI*En_recp/Ac
	Ewald_c[2]=M_PI*En_recp/Ac

#Now compute the remaining terms to complete the Ewald summation

	En_rest=0.0

	for i1 in range(len(x1)):
		for i2 in range(len(x1)):
			del_z=alpha*abs(z1[i1]+z1[i2])
			En_rest+=q1[i1]*q1[i2]*(math.exp(-del_z*del_z)/alpha+math.sqrt(M_PI)*del_z*(erf(del_z)-1)/alpha)
				
	En_rest*=math.sqrt(M_PI)/Ac

	Ewald_c[0]+=-En_rest
	Ewald_c[3]=-En_rest
	
	return Ewald_c