#! /usr/bin/env python

#################################################################################
# Generate atom positions for a (X X X) surface of a monoatomic crystal.
# Print output in Cartesian and Direct coordinates (equilibrium positions).
# Assign initial displacements and velocities to the surface atoms according to 
# independent HO model, for a certain temperature. 
#
# F. Nattino
# June 2013, Leiden

import sys
import string
import math
from random import gauss
from numpy import linalg, dot

#################################################################################
# Define Input Here ############################################################# 

t_exp 	= 1.0014			# thermal expantion coefficient
					# from http://www.ingentaconnect.com/content/matthey/pmr/1997/00000041/00000001/art00007
A 	= 4.092234555 * t_exp  	# expanded bulk lattice constant

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

Zinterlayer = [ 2.3763*t_exp, 2.341*t_exp, 2.3750*t_exp ]	    # Interlayer distances (Angstrom)

#Based on a 3x3 cell, but they are still based on the values obtained with a wrong interlayer spacing
Frequency =  [ [ 9.736633, 9.745601, 7.719476 ], \
	        # omega X ,  omega Y  ,   omega Z  -- Layer 1 (meV)
               [ 11.118959, 11.089457, 10.945031 ], \
	        # omega X ,  omega Y  ,   omega Z  -- Layer 2 (meV)
               [ 11.096719, 11.078720, 10.903800 ] ]
                # omega X ,  omega Y  ,   omega Z  -- Layer 3 (meV)

Packing = 'fcc'    # Packing structure (bcc, fcc, hcp)
Surface = '111'    # Surface (Miller indeces)
MassAmu =  196.966 # Atomic mass (amu)
Vacuum  =  15.	   # Vacuum space (Angstrom)
 
Temperature = 170  # Surface temperature (Kelvin)
Temperature=0

# End of Input ##################################################################
#################################################################################

# Few constants 
AmuToKg = 1.66054e-27
BoltzmannConstant = 1.38065e-23
eVToJoule = 1.60218e-19
DiracConstant = 1.054571e-34

# Functions #####################################################################

def CartesianToDirect( Cart, a1, a2, a3 ):
# Transform cartesian coordinates to direct
	T = [ a1, a2, a3]
	T_ = linalg.inv(T)
	Dir = dot(Cart, T_)

	return Dir 

# Main ##########################################################################

# Print Input
print " \n Lattice constant:  %15.13f " %A
print " Generate slab geometry for a ", N[0], " x ", N[1], " unit cell"
print " Interlayer distances: ", Zinterlayer
print " Vacuum space: %4.3f " %Vacuum
print " Packing structure of the metal: ", Packing
print " Surface : ", Surface

# Define primitive surface unit cell

if Packing == 'fcc' : 

	if Surface == '111' :
		v1 = [   A / math.sqrt( 2. )         , 0.                           ]
		v2 = [   A / ( 2. * math.sqrt( 2. ) ), A /  2. * math.sqrt( 3./2. ) ]
	else :
		print " No data for ", Packing, Surface, " surface. Abort. "
		quit()
elif Packing == 'bcc' :
	
	if Surface == '110' :
                v1 = [   A * math.sqrt( 2. ) / 2.    , - A /  2.                    ]
                v2 = [   A * math.sqrt( 2. ) / 2.    ,   A /  2. 		    ]

	elif Surface == '100' :
		v1 = [   A                           , 0.                           ]
                v2 = [   0.                          , A                            ]

        else :
                print " No data for ", Packing, Surface, " surface. Abort. "
                quit()

else :
	print " No data for ", Packing, " packing. Abort. "
        quit()
	
#v2 = [ - 3.* A / ( 2. * math.sqrt( 2. ) ), 3.*A /  2. * math.sqrt( 3./2. ) ]

# Define supercell
a1 = [ N[0] * v1[0]  , N[0] * v1[1] , 0.                           ]
a2 = [ N[1] * v2[0]  , N[1] * v2[1] , 0.                           ]
a3 = [ 0.            , 0.           , sum( Zinterlayer ) + Vacuum  ]

print " \n Primitive Surface Unit cell: \n "
print " % 15.13f % 15.13f     " %tuple(v1)
print " % 15.13f % 15.13f  \n " %tuple(v2)

print " \n Supercell: \n "
print " % 15.13f % 15.13f % 15.13f    " %tuple(a1)
print " % 15.13f % 15.13f % 15.13f    " %tuple(a2)
print " % 15.13f % 15.13f % 15.13f \n " %tuple(a3)

# List of the translation for atoms in a layer
TranslationsList = []
for i in range( N[0] ):
	for j in range( N[1] ):
		TranslationsList.append( ( float(i), float(j) ) )

# List of the relative translations between layers
LayerShiftList = []
if Packing == 'fcc' :
	if Surface == '111' :
		LayerShiftList.append( ( 0.         , 0.               )) # Shift of the first layer with respect to first
		LayerShiftList.append( ( v1[0]      , v2[1] * 2. / 3.  )) # Shift of the second layer with respect to first
		LayerShiftList.append( ( v1[0] / 2. , v2[1] / 3.       )) # Shift of the third layer with respect to first
		ModuleElement = 3

elif Packing == 'bcc' :
        if Surface == '110' :
		LayerShiftList.append( ( 0.         , 0.               )) 
                LayerShiftList.append( ( v1[0]      , 0.               )) 
		ModuleElement = 2
        if Surface == '100' :
                LayerShiftList.append( ( 0.         , 0.               ))
                LayerShiftList.append( ( v1[0] / 2. , v2[1] / 2.       ))  
                ModuleElement = 2

# Compute values for Surface temperature Displacement

StDevVel = math.sqrt( BoltzmannConstant * Temperature /  ( MassAmu * AmuToKg ) )

StDevPos = []

for i  in range( len( Frequency ) ):
	StDevPosLayer = []
	for j in range(3):
		Omega = Frequency[i][j] * eVToJoule * 1.e-3 / DiracConstant
		StDevPosLayerQ = math.sqrt( BoltzmannConstant * Temperature / ( MassAmu * AmuToKg * Omega **2.)  )
		StDevPosLayer.append( StDevPosLayerQ )
	StDevPos.append( StDevPosLayer )

print " \n Standard deviation for generation of velocities: % 15.13f \n" %StDevVel
for i  in range( len( Frequency ) ):
	print " Standard deviation ( m ) for generation of displacements (X,Y,Z) for layer ", i, " : " 
	print " %.6e %.6e %.6e " %tuple( StDevPos[i])


# Print Cartesian Coordinates

Coordinates = [ 0., 0., 0. ]
Cartesian = []

print " \n Equilibrium Cartesian Coordinates ( Ang ): \n "

for i  in range( len( Zinterlayer ) + 1 ):
	j = i % ModuleElement
	LayerShift = LayerShiftList[j]
	if i != 0 :
		Coordinates[2] = math.fabs( Coordinates[2] + a3[2] - Zinterlayer[ i -1 ] ) % a3[2]
	for Translation in TranslationsList :
		Position = [ 0. , 0. , 0. ] 
		Position[0] = Coordinates[0] + LayerShift[0] + ( Translation[0]* v1[0] + Translation[1] * v2[0] )
		Position[1] = Coordinates[1] + LayerShift[1] + ( Translation[0]* v1[1] + Translation[1] * v2[1] )
		Position[2] = Coordinates[2] 
		Cartesian.append( Position )
		print " % 15.13f % 15.13f % 15.13f  " %tuple(Position)

# Print Direct Coordinates

Direct = []

print " \n Equilibrium Direct Coordinates: \n "

for Atom  in Cartesian :
	Position = CartesianToDirect( Atom, a1, a2, a3 )
	Direct.append( Position )
	print " % 15.13f % 15.13f % 15.13f  " %tuple( Position )

print " \n "


# test ###########################################################
#out = open('CoordinateTest.dat', 'w')
#for i in range(len(Cartesian)):
#	out.write(' % 15.13f % 15.13f % 15.13f	 % 15.13f % 15.13f % 15.13f\n' % (Cartesian[i][0], Cartesian[i][1], Cartesian[i][2], Direct[i][0], Direct[i][1], Direct[i][2]))
#out.close()
#################################################################

# Print Cartesian Coordinates with temperature 

print " \n Cartesian Coordinates with temperature ( Ang ): \n "

Coordinates = [ 0., 0., 0. ]
Potential = 0.

for i  in range( len( Cartesian ) ):
	LayerNumber =  i // ( N[0] * N[1] )
 	if  LayerNumber <= ( len( Frequency ) - 1 ):
		for j in range(3):
			dq = gauss( 0., StDevPos[LayerNumber][j] ) 
			Potential = Potential + 0.5 * dq **2. * MassAmu * AmuToKg * ( Frequency[LayerNumber][j] * eVToJoule * 1.e-3 / DiracConstant ) **2.
			Coordinates[j] = Cartesian[i][j] + dq * 1.e10 
		print " % 15.13f % 15.13f % 15.13f  " %tuple( Coordinates )
	else:
		print " % 15.13f % 15.13f % 15.13f  " %tuple( Cartesian[i] )


print " \n Velocities ( Ang / fs ): \n "

# Print Velocities
Velocities = [ 0. , 0., 0. ]
Kinetic = 0.

for i  in range( len( Cartesian ) ):
        LayerNumber =  i // ( N[0] * N[1] )
        if  LayerNumber <= ( len( Frequency ) - 1 ):
                for j in range(3):
                        dVq = gauss( 0., StDevVel ) 
			Kinetic = Kinetic + 0.5 * dVq **2. * MassAmu * AmuToKg 
                        Velocities[j] = dVq * 1.e-5
                print " % 15.13f % 15.13f % 15.13f  " %tuple( Velocities )
        else:
                print " % 15.13f % 15.13f % 15.13f  " %tuple( [0.,0.,0.] )


print " \n  Potential computed for independent HO ( K ):		%15.10f" %( Potential / ( 0.5 * len( Frequency ) * N[0] * N[1] * 3. * BoltzmannConstant ))
print "  Kinetic Energy computed for independent HO ( K ):	%15.10f\n" %( Kinetic / ( 0.5 * len( Frequency ) * N[0] * N[1] * 3. * BoltzmannConstant ))