#! /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

# List of Modifications by D. Migliorini
# 09/01/2017 	- Adapted to handle an fcc(211) surface.
#		- Added Boltzmann theory in header
# 
# the velocities are assigned randomly according to a independent HO model as 
# explained in http://hyperphysics.phy-astr.gsu.edu/hbase/Kinetic/maxspe.html
# the distribution for the position can be derived similarly starting from a 
# Boltzmann distribution Ae**-(E/kT) considering the HO potential energy
# U = 0.5 * mass * ( omega*q )**2
# 

import sys
import string
import math
from random import gauss
from numpy import linalg, dot,array
from numpy.linalg import inv
from operator import itemgetter
import vasptools as vt


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

t_exp 	= 1.005311542			# thermal expantion coefficient at 650K estimated as average between 600 K and 700 K
					# from http://www.ingentaconnect.com/content/matthey/pmr/1997/00000041/00000001/art00007
units 		= [1, 3]
layers		=  4
frozen_layers 	= 1
 
Frequency =  	[ [ 11.619459, 10.976146, 12.324700 ] , [ 14.461519, 12.881512, 9.7349000 ] , [ 14.400451, 11.967824, 8.1527640 ],   
# omega X,Y,Z	  step edge atoms			central row atoms		    third row atoms			  -- Layer 1 (meV)
		  [ 12.395621, 14.290262, 15.265980 ] , [ 12.269282, 13.502896, 13.630143 ] , [ 14.427735, 13.242476, 12.265454 ],
# omega X,Y,Z	  step edge atoms			central row atoms		    third row atoms			  -- Layer 2 (meV)
		  [ 13.466998, 14.115274, 14.482186 ] , [ 13.287758, 13.383251, 13.087410 ] , [ 13.552076, 13.419181, 12.946799 ] ]
# omega X,Y,Z	  step edge atoms			central row atoms		    third row atoms			  -- Layer 3 (meV)

MassAmu 	=  195.080 	# Atomic mass (amu)
VACUUM  	=   13.		# Vacuum space (Angstrom)
Temperature 	=  650		# Surface temperature (Kelvin)

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

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

# Ideal Slab preparation ########################################################
################################################################################# 

slab = '/home/miglio/CHD3_Pt211/slab_optg/SRP032-DF/ALL/SLABS_5Ang/surface_%dl_1x1_5v' % layers

T  = vt.get_cell(slab)
positions = array(vt.get_positions(slab))
Natoms  = len(positions)

# Store the equilibrium positions as cartesian ------------------------------
positions_cart = dot(positions, T)

# Resize the vacuum ---------------------------------------------------------
T[2][2] += VACUUM -5.
T = array(T)
T_ = inv(T)

# Updated direct positions considering new vacuum ---------------------------
positions = dot(positions_cart, T_)

# Unit Cell Lattice Vector (before thermal expantion)------------------------
print "\n Unit Cell Lattice Vector (before thermal expantion) "
for i in range(3): print " % 15.13f % 15.13f % 15.13f  " % tuple(T[i])

print "\n CARTESIAN POSITION BEFORE EXPANTION"
for i in range(Natoms): print " % 15.13f % 15.13f % 15.13f  " % tuple(dot(positions[i], T))

# Termal Expantion of the cell ----------------------------------------------
for i in range(3): T[i] = T[i]*t_exp
T_ = inv(T)

# Unit Cell Lattice Vector (after thermal expantion)-------------------------
print "\n Unit Cell Lattice Vector - %d K; expantion:  %10.8f:" % (Temperature, t_exp)
for i in range(3): print " % 15.13f % 15.13f % 15.13f  " % tuple(T[i])

print "\n CARTESIAN POSITION AFTER EXPANTION"
for i in range(Natoms): print " % 15.13f % 15.13f % 15.13f  " % tuple(dot(positions[i], T))

# ----------------------------------------------------------------------------
# Generate SuperCell ---------------------------------------------------------

newT, newpositions = vt.repeat(positions, [units[0],units[1],1], T)     # generate the supercell
newT_ = inv(newT)

atomsXlayer     = 3*units[0]    *       units[1]                # compute the slab features
surf_atoms      = atomsXlayer   *       layers                  #
free_layers     = layers        -       frozen_layers           #
free_atoms      = free_layers   *       atomsXlayer             #
frozen_atoms    = frozen_layers *       atomsXlayer             #
Natoms		= len(newpositions)
newpositions = sorted(newpositions, key = itemgetter(2), reverse = True)

# ----------------------------------------------------------------------------
# Main ##########################################################################
'''
print " \n Generate Supercell\n [ %d x %d x 1 ] - %d K; expantion:  %10.8f:" % (units[0], units[1], Temperature, t_exp)
for i in range(3): print " % 15.13f % 15.13f % 15.13f  " % tuple(newT[i])
print " Vacuum space: %6.3f " % (VACUUM*t_exp)
'''

newcartesian = dot(newpositions, newT)
'''print "\n SUPERCELL CARTESIAN POSITION AFTER EXPANTION"
for i in range(Natoms): print " % 15.13f % 15.13f % 15.13f  " % tuple(newcartesian[i])'''

#################################################################################
# Compute Values for Surface Temperature Displacements ##########################
#################################################################################

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//3, " : " 
	print " %.6e %.6e %.6e " %tuple( StDevPos[i])


#for i in range(len(Frequency)):
#	print " % 15.13f % 15.13f % 15.13f  |  % 15.13f % 15.13f % 15.13f  " % (Frequency[i][0], Frequency[i][1], Frequency[i][2], StDevPos[i][0], StDevPos[i][1], StDevPos[i][2])


# Print Coordinates #############################################################

print " \n Supercell (%d K):\n" % Temperature
for i in range(3): print " % 15.13f % 15.13f % 15.13f  " % tuple(newT[i])

print " \n Equilibrium Cartesian Coordinates ( Ang ): \n "
for i in range(Natoms):	print " % 15.13f % 15.13f % 15.13f  " % tuple(newcartesian[i])

print " \n Equilibrium Direct Coordinates: \n "
for i in range(Natoms): print " % 15.13f % 15.13f % 15.13f  " % tuple(newpositions[i])
print " \n "


# Print Cartesian Coordinates with temperature 

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

'''
print "\n Atom Position Test:"
for i in range( Natoms ): 
	layer   = i // atomsXlayer
	row     = ( i  - atomsXlayer*(layer)) //3
	print " Atom %d is in layer %d and row %d and with freq %d " % ( i, layer, row, layer*3+row )
'''
Potential = 0.

for i in range( Natoms ):

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

	layer 	= i // atomsXlayer 
	row 	= ( i  - atomsXlayer*(layer)) //3 
	if layer+1  <= free_layers:
		for j in range(3):
			freq = layer*3+row
			dq = gauss( 0., StDevPos[freq][j] )
			Potential = Potential + 0.5 * dq **2. * MassAmu * AmuToKg * ( Frequency[freq][j] * eVToJoule * 1.e-3 / DiracConstant ) **2.
			Coordinates[j] = newcartesian[i][j] + dq * 1.e10
#			print StDevPos[freq][j], math.sqrt( BoltzmannConstant * Temperature / ( MassAmu * AmuToKg * ( Frequency[freq][j] * eVToJoule * 1.e-3 / DiracConstant ) **2.)  ) 
		print " % 15.13f % 15.13f % 15.13f  " %tuple( Coordinates )	
	else:
		print " % 15.13f % 15.13f % 15.13f  " %tuple( newcartesian[i] )

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

# Print Velocities
Kinetic = 0.

for i in range( Natoms ):

	Velocities = [ 0. , 0., 0. ]
        layer   = i // atomsXlayer
        if layer+1  <= free_layers:
                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 "\nfree_layers: ", free_layers, "atoms X layer: ", atomsXlayer

print " \n  Potential computed for independent HO ( K ):		%15.10f" %( Potential / ( 0.5 * free_layers * atomsXlayer * 3 *  BoltzmannConstant ))
print "  Kinetic Energy computed for independent HO ( K ):	%15.10f\n" %( Kinetic / ( 0.5 * free_layers*atomsXlayer * 3 *  BoltzmannConstant ))