#!/usr/bin/env python

# The idea of this scrpt is to check the temperature (T) and 
# the interlayer distances (d) of a bare metal slab in 
# a dynamic reading the VASP XDATCAR file and return
# T vs time, < T >, d vs time and < d > 


# ===================================================== #
from 	os 		import chdir, remove, path
from 	numpy 		import dot, sqrt, array, mean
from 	numpy.linalg 	import inv, norm
from 	re			import match 
import	commands
# ===================================================== #

au2Ang	=	0.52917721092				# some conversion constants
amu2au	=	1822.88839				#
fs2au	=	41.341373337				#
au2eV	=	27.21138505				#
deg2rad	=	0.0174532925				#
amu2kg	=	1.66053892 	* 10**-27		#
J2eV	=	6.24181		* 10**18		#
h	=	6.6260695729    * 10**-34		#
BoltzmannConstant = 1.38065e-23				# [(m**2) * (kg)] / [(s**2) * (K)]

# ===================================================== #

T_ideal   	= 170		# K
Nfree	  	= 27		# moving atoms
Natoms		= 38		# total number of atoms
Nsurfatoms	= 36		# Number of surface atoms
tstep	  	= 1.0		# fs
MassAtom  	=196.966	# amu
AtomsXlayer 	= 9

# ===================================================== #
def read_xdatcar(infile):

	lines = open(infile+'/states.xyz').readlines()

	### cell parameters ###
	POSlines = open(infile+'/POSCAR').readlines()           # read file

	a = [float(x) for x in  POSlines[2].split()]                   # store cell parameters
	b = [float(x) for x in  POSlines[3].split()]                   #
	c = [float(x) for x in  POSlines[4].split()]                   #
	Cell = array( [a, b, c] )
	Cell_ = inv( Cell )
	
	### read positions ###

	cart = []
        Nsteps = int(lines[-(Nsurfatoms+3)].split()[2])
	header  = 1
	block   = Natoms + 2 # surf atoms + HCl + headline
        for step in range(Nsteps):
		appo = []
		for atom in range(Nsurfatoms):
			appo.append( [ float(x) for x in lines[( header + step*block + atom + 1)].split()[1:4] ] )
		cart.append(array(appo))

	cart = array(cart)

	### convert to direct coordinates

        direct = []
        for step in range(Nsteps):
                appo = []
                for at in range(Nsurfatoms):
                        appo.append( dot(cart[step][at], Cell_))
                direct.append(array(appo))

        direct = array(direct)


	### consider boundaries ####

	diff = [0.,0.,0.]
	real = []
	appo = []
	for i in range(Nsurfatoms): 
		appo.append( [ float(x) for x in direct[0][i] ] )
	real.append(appo)

        for step in range(1, Nsteps):
		appo = []
                for at in range(Nsurfatoms):
                        appo.append(direct[step][at])
			for j in range(3):
				diff[j] = appo[at][j] - real[step-1][at][j]
				if   diff[j] < -0.5 :
					appo[at][j] += 1.
				elif diff[j] >  0.5  :
					appo[at][j] -= 1.
		real.append(appo)
	real = array(real)	

        ### direct2cartesian ###

	cartesian = []
	for step in range(Nsteps): 
		appo = []
		for at in range(Nsurfatoms):		
			appo.append( dot(real[step][at], Cell))
		cartesian.append(array(appo))

	cartesian = array(cartesian)
	#######################
        ### boudary test ###
#        direct_test     = open('direct.dat', 'w')
#        real_test       = open('real.dat', 'w')
#        cart_test       = open('cartesian.dat', 'w')
#
#        for step in range(Nsteps):
#                direct_test.write('%11.8f %11.8f\n' % tuple(direct[step][0][0:2]))
#                real_test.write('%11.8f %11.8f\n' % tuple(real[step][0][0:2]))
#                cart_test.write('%11.8f %11.8f\n' % tuple(cartesian[step][0][0:2]))
#        direct_test.close()
#        real_test.close()
#        cart_test.close()

	return  cartesian, real, Cell, Nsteps, Nsurfatoms

# ====================================================== #

def slab_analysis(mypath, w=False):

	cartesian, direct, Tcell, Nsteps, Natoms = read_xdatcar(mypath)

	### velocities ############

	Vc 	= []
	V  	= [0.,0.,0.]
	appo 	= []
	
	for at in range(Natoms):
                        appo.append([99.99999999,99.99999999,99.99999999])
	Vc.append(appo)

	for step in range(1, Nsteps-1):
		appo = []
		for at in range(Natoms):
			for j in range(3):
				V[j] = float( cartesian[(step+1)][at][j] - cartesian[step-1][at][j]  ) / (2.*tstep)
			appo.append([x for x in V])
		Vc.append(appo)

	appo = []
	for at in range(Natoms):
		appo.append([99.99999999,99.99999999,99.99999999])
	Vc.append(appo)

	Vc = array(Vc)

#        v_test          = open('v.dat', 'w')
#        for step in range(Nsteps):
#                v_test.write('%11.8f %11.8f\n' % tuple(Vc[step][6][0:2]))
#        v_test.close()

	# since the velocities in the poscar and in the contcar file are give 1/2*dt off in respect to the positions
	# it's not possible to calculate them for the first and the last step on the XDATCAR coordiantes
	# therefore all the v of these two steps have been setted to 99.99999999999 

	### check T and kinetic ###

	kinetic = []
	T 	= []
	for step in range(1, Nsteps-1):
		Ek = 0.0
		for at in range(Natoms):  
			Vtot = norm( Vc[step][at] )
			Ek += 0.5 * MassAtom * Vtot**2  * ( amu2au * 1./fs2au**2 * 1./au2Ang**2 )
		kinetic.append( Ek*au2eV )

		step_temp = 2 * kinetic[-1] / ( J2eV * BoltzmannConstant * 3 * Nfree ) 
		T.append(step_temp)

  
	# Test T and Ek #########################################################

	v_test          = open('T_Ek.dat', 'w')
	for step in range(len(T)):
		v_test.write('%11.8f %11.8f\n' % (T[step], kinetic[step]))
	v_test.close()

	#########################################################################

	Tavg = mean(T)
	
	StdDevT_NUM = 0.

	for i in range(len(T)): 	StdDevT_NUM += ( T[i]-Tavg )**2 
	StdDevT = sqrt( StdDevT_NUM / len(T) )	

	### interlayer distances ###

	layers 		= Natoms / AtomsXlayer 
	interlayers 	= layers -1
	Z		= []
	for step in range(Nsteps):
		appo = []
		for layer in range(layers):
			avgZ = 0
			for at in range(AtomsXlayer*layer, AtomsXlayer*(layer+1)):
				if layer == 0 and cartesian[step][at][2] < Tcell[2][2]/2:
					avgZ += (cartesian[step][at][2] + Tcell[2][2]) / AtomsXlayer
				else:
					avgZ += (cartesian[step][at][2]) / AtomsXlayer
			appo.append(avgZ)
		Z.append(appo)

	Z = array(Z)

	dZ = []
	for step in range(Nsteps):
		appo = []
		for inter in range(interlayers):
			appo.append(Z[step][inter] - Z[step][inter+1])
		dZ.append(appo)
	dZ = array(dZ)

	dZavg = []

	for inter in range(interlayers):
		appo = 0.0
		for step in range(Nsteps): appo += dZ[step][inter] / (Nsteps-1)
		dZavg.append(appo)

	### printing results ###


	if path.exists(mypath+'/T_check.dat'): remove(mypath+'/T_check.dat')
	if path.exists(mypath+'/layers_check.dat'): remove(mypath+'/layers_check.dat')

	out1 = open(mypath+'/T_check.dat', 'w')
	out2 = open(mypath+'/layers_check.dat', 'w')
	for step in range(Nsteps):
#		out2.write(' %d	%.8f	%.8f	%.8f	%.8f\n' % (step, dZ[step][0], dZ[step][1], dZ[step][2], dZ[step][3]))
		out2.write(' %d %.8f	%.8f	%.8f\n' % (step, dZ[step][0], dZ[step][1], dZ[step][2]))
	for step in range(Nsteps-2):
		out1.write(' %d	%.8f	%.8f\n' % (step+1, Tavg, T[step]))
	out1.close()
	out2.close()

        ### print ###################
	if w==True:
		n = int(commands.getstatusoutput('ls /home/nick/au111-HCl-NN-MD/SurfaceEquilibration/EQUILIBRATED/ | wc -l' )[1] ) + 1
		print "surface", n
		output = open('/home/nick/au111-HCl-NN-MD/SurfaceEquilibration/EQUILIBRATED/AIMDSurf%.2dAuCarAndVel_3x3_4L_MSRPBE-170K.dat' % n, 'w')

		for step in range(1,Nsteps-1):
			output.write(' %d\n' % (step+1))
			for atom in range(Natoms):
				output.write('  %11.8f %11.8f %11.8f\n' % tuple( cartesian[step][atom] ))
			for atom in range(Natoms):
				output.write('  %11.8f %11.8f %11.8f\n' % tuple(Vc[step][atom]))
		output.close()

	# P R I N T ################################################################

	print 20*'*', 'SLAB EQUILIBRATION ANALYSIS', 19*'*', '\n'

	print 'check steps:  ', Nsteps, len(Vc), len(cartesian)
#	print 'Standard Deviation Numerator:	%.1f\n' %  StdDevT_NUM

	print 'Average interlayer distances:'
	print '<dZ>	=',
	for i in range(interlayers): print ' %.4f ' % dZavg[i], 
	print '\n\ninitial T	= %.1f K' % T[0]
	print '<T>		= %.1f +/- %.1f K\n'% (Tavg, StdDevT)

	print 68*'*'

	return T, StdDevT_NUM

#####################################################################################################################################
from sys import argv

ShouldIWrite = False

if len(argv) == 1: 
	appoT, appo_NUM = slab_analysis('.', ShouldIWrite)
else: 
	NUM  = 0.
	Tavg = 0.
	steps= 0.
	all_avgT = []
	for i in range(1, len(argv)):
		print 'SLAB ', i, 60*'*' 
		appoT, appo_NUM = slab_analysis(argv[i], ShouldIWrite)
		NUM 	+= appo_NUM		
		steps 	+= len(appoT)
		Tavg	+= sum(appoT)
		all_avgT.append(mean(appoT))


	Tavg = Tavg/steps
	StdDev = sqrt(NUM/steps)

	
	t_RANGE = [Tavg-StdDev, Tavg+StdDev]

	print 68*'*', '\n', 68*'*'

	print 'Slab equilibration analyzed:	%d' % (len(argv)-1)
	print 'Slabs\' <T>:'
        for i in range(len(all_avgT)):
                if all_avgT[i] > t_RANGE[0] and all_avgT[i] < t_RANGE[1]:
                        print '\033[92m'+'%.1f' % all_avgT[i],
                else:
                        print '\033[91m'+'%.1f' % all_avgT[i],

	print  '\033[0m'+'\n'
	print 'Ideal T				%.1f K ' % T_ideal
	print 'Total average temperature <T>:	%.1f +\- %.1f K'	% (Tavg, StdDev)
	

	print 68*'*', '\n', 68*'*'