#!/usr/bin/env python3
from ase import Atoms, units
from ase.build import surface, add_vacuum, bulk, molecule, add_adsorbate
from ase.io import read, write
from ase.constraints import FixAtoms, FixedLine
from ase.lattice.cubic import FaceCenteredCubic
from operator import itemgetter
from os import path
from random import gauss
import sys, os
import numpy as np

InitialPOS = open('DynamicsINPUTS/inicond_D2_FWHM/Output/ICON.txt', 'r').readlines()

layers = int(sys.argv[1])
supercellx = int(sys.argv[2])
supercelly = int(sys.argv[3])
mil_h = int(sys.argv[4])
mil_k = int(sys.argv[5])
mil_l = int(sys.argv[6])
if sys.argv[7] == 'True':
	motion = True
elif sys.argv[7] == 'False':
	motion = False
else:
	print('WARNING: something wrong with the input for the surface motion')
	exit(1)
vacuum = 15.
Ts = 0.

if mil_h % 2 == 0 and mil_k % 2 == 0 and mil_l % 2 == 0:
	print('WARNING: miller indices are divisable by 2, and should be')
	exit(1)

# Read the relaxed slab from a file
POSCARname = 'POSCAR_{0:d}{1:d}{2:d}-{3:d}L'.format( mil_h, mil_k, mil_l, layers )
if not path.exists(POSCARname):
	print('WARNING: POSCAR of the slab does not exist')
	exit(1)
else:
	slab = read(POSCARname, format='vasp')

# Remove any constraints that might have been present in the POSCAR that was read
del slab.constraints

# Determine size of the terrace
if mil_h == mil_l and mil_h == mil_k:	# 111 surface without steps
	Nterrace = 1
elif mil_k == mil_l:			# 111 surface with 100 steps
	if (mil_h - mil_l) == 1:
		Nterrace = mil_h*2 - 1
	elif (mil_h - mil_l) == 2:
		Nterrace = mil_h - 1
	else:
		print('WARNING: miller index not implemented for Nterrace')
		exit(1)
elif mil_h == mil_k:			# 111 surface with 110 steps
	if (mil_h - mil_l) == 1:
		Nterrace = mil_h*2
	elif (mil_h - mil_l) == 2:
		Nterrace = mil_h
	else:
		print('WARNING: miller index not implemented for Nterrace')
		exit(1)

# Extend the slab in the x and y direction as requested
slab = slab.repeat((supercellx, supercelly, 1))

def order_atoms( atoms ):
	"""
	This function orders the atoms from top to down.
	Function returns the modified atoms object.
	"""
	
	orders = [(atom.index, round(atom.x, 3), round(atom.y, 3), -round(atom.z, 3), atom.index) for atom in atoms]
	orders.sort(key=itemgetter(3, 1, 2))
	newatoms = atoms.copy()
	for index, order in enumerate(orders):
		newatoms[index].position = atoms[order[0]].position.copy()
	
	return newatoms

# Renumber systematically from top to down
slab = order_atoms( slab )

Cell = slab.get_cell()
# Get rid of some annoying effects of having negative values of Z
for i in range( len( slab ) ):
        if slab[i].position[2] < 0.:
                 slab[i].position[2] += Cell[2][2]
# Translate all atoms in order to get the top step atom at the (0, 0, 0) position
min_idx = np.argmin(slab.positions[:,2])
max_idx = np.argmax(slab.positions[:,2])
slab.translate( -slab.positions[max_idx,:] )
slab.wrap(pbc=(1,1,0))

# Here we add vacuum
Cell[2][2] = slab[max_idx].position[2] - slab[min_idx].position[2]
slab.set_cell(Cell)
add_vacuum(slab, vacuum)

# Renumber systematically from top to down, again.
slab = order_atoms( slab )

# Determine unit cell vectors
Cell = slab.get_cell()
Cell_= np.linalg.inv(Cell)

def surface_motion( atoms, Ts ):
	"""
	Surface atoms are displaced according to a surface temperature based on a harmonic oscillator model
	"""
	
	#Read the frequencies from a file
	FREQname = 'FREQUENCIES_{0:d}{1:d}{2:d}-{3:d}L'.format( mil_h, mil_k, mil_l, layers )
	if not path.exists(FREQname):
		print('WARNING: Frequency list of the slab does not exist')
		exit(1)
	else:
		lines = open(FREQname, 'r').readlines()
	
	# The format in the file assumes that the frequencies correspond to how the slab is ordered.
	# I.e, from top to down.
	Frequency = [ [ float(x) for x in line.split() ] for line in lines ]

	# Few constants 
	AmuToKg = 1.66054e-27
	BoltzmannConstant = 1.38065e-23
	eVToJoule = 1.60218e-19
	DiracConstant = 1.054571e-34
	MassAmu = 195.080	# Atomic mass of surface atom (amu)

	StDevPos = []
	for i in range( len( Frequency ) ):
		StDevPosLayer = []
		for j in range(3):
			Omega = Frequency[i][j] * eVToJoule * 1.e-3 / DiracConstant
			StDevPosLayerQ = np.sqrt( BoltzmannConstant * Ts / ( MassAmu * AmuToKg * Omega **2.)  )
			StDevPosLayer.append( StDevPosLayerQ )
		StDevPos.append( StDevPosLayer )
	
	# Displacements of the surface atoms according to the surface temperature (Ts)
	# using a independent harmonic oscillator model
	#Potential = 0.
	for i in range( Nterrace * supercellx * supercelly * 3 ):
		k = i // (supercellx * supercelly)
		for j in range(3):
			dq = gauss( 0., StDevPos[k][j] )
#			Potential = Potential + 0.5 * dq **2. * MassAmu * AmuToKg * ( Frequency[i][j] * eVToJoule * 1.e-3 / DiracConstant ) **2.
			slab.positions[i][j] += dq * 1.e10
	
	#print( "Potential computed for independent HO ( K ): {:15.10f}".format( Potential / ( 0.5 * len( Frequency ) * supercellx*Nterrace*supercelly * 3. * BoltzmannConstant )) )
	
	return

if motion == True:
	surface_motion( slab, Ts )

def position(Z, r, X, Y, THETA, PHI):
	"""
	Convert spherical coordinates to cartesian coordinates
	"""
	PHI     = np.radians(PHI)
	THETA   = np.radians(THETA)
	COM = np.array( [X, Y, Z] )
	SPH2CAR = np.array( [ r*np.sin(THETA)*np.cos(PHI), r*np.sin(THETA)*np.sin(PHI), r*np.cos(THETA) ] )
	POS = [ COM - 0.5*SPH2CAR, COM + 0.5*SPH2CAR ]

	return POS

H2 = molecule('H2')
add_adsorbate( slab, H2, height=vacuum/2. )

# Here we generate the POSCARs
for i in range(1, len(InitialPOS) + 1):
	if not os.path.exists('{0:06d}'.format( i ) ):
		os.makedirs('{0:06d}'.format( i ) )
	
	X_d = float( InitialPOS[i-1].split()[0] )
	Y_d = float( InitialPOS[i-1].split()[1] )
	COM = np.dot( [X_d, Y_d, 0.], Cell )
	COM[2] = vacuum/2.
	r = float( InitialPOS[i-1].split()[2] )
	theta = float( InitialPOS[i-1].split()[3] )
	phi = float( InitialPOS[i-1].split()[4] )
	slab.positions[-2:] = position(COM[2], r, COM[0], COM[1], theta, phi)
	
	vel = np.zeros( ( len(slab), 3 ) )
	vel[-2] = [ float( InitialPOS[i-1].split()[5] ) / 1000., float( InitialPOS[i-1].split()[6] ) / 1000., float( InitialPOS[i-1].split()[7] ) / 1000. ]
	vel[-1] = [ float( InitialPOS[i-1].split()[8] ) / 1000., float( InitialPOS[i-1].split()[9] ) / 1000., float( InitialPOS[i-1].split()[10] ) / 1000. ]
	vel /= units.fs
	slab.set_velocities( vel )
	
	write('{0:06d}/cfg_lammps.RuN.config'.format( i ), slab, format='lammps-data', velocities=True)
	#write('{0:06d}/POSCAR'.format( i ), slab, format='vasp')
	#write('POSCAR', slab, format='vasp', direct='true')