#!/usr/bin/env python

import sys, os
from math import *
from ReadTrajectoryVASP import *
from MakeSlab import Slab 

if (__name__ == "__main__"):

	MassWamu = 183.850
        amu2au = 1822.88839
        Angstrom2au = 1.889725989
        fs2au = 41.3413732914

        au2K = 315775.04

	A =  4.01679688 * 1.000485  # Bulk lattice constant

	N = [ 0. , 0. ]
	N[0] = 3                            # Number of atoms along a1
	N[1] = 3                            # Number of atoms along a2

	Zinterlayer = [ 2.342135, 2.3011155, 2.3011155, 2.342135 ]             # Interlayer distances (Angstrom)

	Packing = 'fcc'    # Packing structure (bcc, fcc, hcp)
	Surface = '111'    # Surface (Miller indeces)
	Vacuum  = 13.      # Vacuum space (Angstrom)	

	MovingAtoms = 36
	NAtomsPerLayer = N[0] * N[1] 
	FirstAtomThirdLayer = 2 * NAtomsPerLayer
	LastAtomThirdLayer  = 3 * NAtomsPerLayer


	TheSlab = Slab( Packing, Surface, A, Zinterlayer, N, Vacuum)

	Basis  = TheSlab.GenerateBasisVectors( )
	rAtoms = TheSlab.GenerateCoordinates( )

	del TheSlab

	TimeStep = 1.0 # fs 

	for surf in range( 1, 11):

		Traj = Trajectory( Verbosity=0 ) 
                Directory = 'traj%02i' %surf + '/'
	
       	        File = Directory + 'POSCAR_3'
                Traj.ReadPOSCAR( File )
                File = Directory + 'XDATCAR_3'
                Traj.ReadXDATCAR(Filename = File)
                Traj.ExtractRealPositionsAndVelocities( TimeStep )
                File = Directory + 'CONTCAR'
                Traj.ReadCONTCAR(Filename = File) 

		AveragePositions   = [ ]
		DisplacementsEq    = [ ]
		DisplacementsTime0 = [ ]
	
		RMS3D = [ ]
		RMSX  = [ ]
		RMSY  = [ ]
		RMSZ  = [ ]

                for i in range( 1, Traj.maxnt ):
			for j in range( MovingAtoms ):
                               if i == 1:
                                        AveragePositions.append( [ 0. , 0. , 0. ] )
                               for k in range( 3 ):
                                        AveragePositions[ j ][ k ]   = AveragePositions[ j ][ k ]   + Traj.rCar[i][j][k] / float( Traj.maxnt - 1 )

		for i in range( 1, Traj.maxnt ):

			DispXYZ = []
			Disp3D  = []

			for j in range( MovingAtoms ):

				if i == 1:
					rAtoms[j] = Traj.FindPointClosestToImage( i, rAtoms[j], j ) 

#				DisplacementX = Traj.rCar[i][j][0] - rAtoms[j][0]
#				DisplacementY = Traj.rCar[i][j][1] - rAtoms[j][1]
#				DisplacementZ = Traj.rCar[i][j][2] - rAtoms[j][2]

#                                DisplacementX = Traj.rCar[i][j][0] - Traj.rCar[1][j][0] 
#                                DisplacementY = Traj.rCar[i][j][1] - Traj.rCar[1][j][1]
#                                DisplacementZ = Traj.rCar[i][j][2] - Traj.rCar[1][j][2]

				DisplacementX = Traj.rCar[i][j][0] -  AveragePositions[j][0]
                                DisplacementY = Traj.rCar[i][j][1] -  AveragePositions[j][1]
                                DisplacementZ = Traj.rCar[i][j][2] -  AveragePositions[j][2]

				Displacement3D = sqrt( DisplacementX ** 2. + DisplacementY ** 2. + DisplacementZ ** 2. )

				if i == 1:
					RMS3D.append( 0. )
					RMSX.append( 0. )
					RMSY.append( 0. )
					RMSZ.append( 0. )

				RMS3D[ j ] = RMS3D[ j ] + Displacement3D ** 2. / float( Traj.maxnt - 1 )
				RMSX[ j ] = RMSX[ j ] + DisplacementX ** 2. / float( Traj.maxnt - 1 )
				RMSY[ j ] = RMSY[ j ] + DisplacementY ** 2. / float( Traj.maxnt - 1 )
				RMSZ[ j ] = RMSZ[ j ] + DisplacementZ ** 2. / float( Traj.maxnt - 1 )

				DispXYZ.append( DisplacementX )
				DispXYZ.append( DisplacementY )
				DispXYZ.append( DisplacementZ )				

				Disp3D.append( Displacement3D )
			print i , len(Disp3D) * "% 10.7f " %tuple ( Disp3D )

                del Traj

		print "\n"

		RMSAllAtoms   = 0.

		RMSFirstLayer = 0.

		RMSXFirstLayer = 0.
		RMSYFirstLayer = 0.
		RMSZFirstLayer = 0.

		RMSThirdLayer = 0.

                RMSXThirdLayer = 0.
                RMSYThirdLayer = 0.
                RMSZThirdLayer = 0.


		for j in range( MovingAtoms ):

			#RMS3D[ j ] = sqrt( RMS3D[ j ] )
                        #RMSX[ j ]  = sqrt( RMSX[ j ] )
                        #RMSY[ j ]  = sqrt( RMSY[ j ] )
                        #RMSZ[ j ]  = sqrt( RMSZ[ j ] )



			# Root mean square displacement averaged over all atoms
			RMSAllAtoms = RMSAllAtoms + RMS3D[ j ] / float( MovingAtoms )

			# Root mean square displacement averaged over atoms of the first layer
			if j in range( NAtomsPerLayer ):
				# root mean square 3D displacement
				RMSFirstLayer = RMSFirstLayer + RMS3D[ j ] / float( NAtomsPerLayer )
				# root mean square X, Y, Z displacements
				RMSXFirstLayer = RMSXFirstLayer + RMSX[ j ] / float( NAtomsPerLayer )
				RMSYFirstLayer = RMSYFirstLayer + RMSY[ j ] / float( NAtomsPerLayer )
				RMSZFirstLayer = RMSZFirstLayer + RMSZ[ j ] / float( NAtomsPerLayer )

			# Root mean square displacement averaged over atoms of the third layer (bulk like)
			if j in range( FirstAtomThirdLayer, LastAtomThirdLayer ) :
				RMSThirdLayer = RMSThirdLayer + RMS3D[ j ] / float( NAtomsPerLayer )
                                # root mean square X, Y, Z displacements
                                RMSXThirdLayer = RMSXThirdLayer + RMSX[ j ] / float( NAtomsPerLayer )
                                RMSYThirdLayer = RMSYThirdLayer + RMSY[ j ] / float( NAtomsPerLayer )
                                RMSZThirdLayer = RMSZThirdLayer + RMSZ[ j ] / float( NAtomsPerLayer )


		print 100*"#"
		print "# Surface ", surf 
		print "# RMS Displacements [ sqrt( < u ^ 2 >_t ) ] - Averages: "
		print "# -> Over all atoms       = ", sqrt(RMSAllAtoms)
		print "# -> Over atoms I   layer = ", sqrt(RMSFirstLayer)
		print "# -> Over atoms III layer = ", sqrt(RMSThirdLayer)
		print "#"
		print "# RMS Displacement in X, Y, Z [ sqrt( < u_X ^ 2 >_t ) ] - Average over I   Layer: "
		print "# -> X: ", sqrt(RMSXFirstLayer)
		print "# -> Y: ", sqrt(RMSYFirstLayer)
		print "# -> Z: ", sqrt(RMSZFirstLayer)
                print "# RMS Displacement in X, Y, Z [ sqrt( < u_X ^ 2 >_t ) ] - Average over III Layer: "
                print "# -> X: ", sqrt(RMSXThirdLayer)
                print "# -> Y: ", sqrt(RMSYThirdLayer)
                print "# -> Z: ", sqrt(RMSZThirdLayer)
		print 100*"#"