#!/usr/bin/env python
from ase import Atoms
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 random import gauss
import sys
import numpy as np

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])
Zmin = float(sys.argv[7])
Zmax = float(sys.argv[8])
rmin = float(sys.argv[9])
rmax = float(sys.argv[10])
vacuum = 15.
lattice = 4.005185133 #PBEa57-vdW-DF2
bulk = bulk('Pt', crystalstructure='fcc', a=lattice)

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)

# Determine size of the terrace and direction in which the crystal is cut
if mil_h == mil_l and mil_h == mil_k:	# 111 surface without steps
	Nterrace = 1
	single_divider = 3
	directions=[[1, -1, 0],[1, 0, -1],[1, 1, 1]]
elif mil_k == mil_l:			# 111 surface with 100 steps
	if (mil_h - mil_l) == 1:
		Nterrace = mil_h*2 - 1
		directions=[[mil_l*4, -mil_h*2, -mil_h*2], [0, 1, -1], [mil_h, mil_k, mil_l]]
		single_divider = int( round( 3*mil_h - 2 ) )
		if single_divider % 2 == 0:
			single_divider = int( round(single_divider/2) )
	elif (mil_h - mil_l) == 2:
		Nterrace = mil_h - 1
		directions=[[(mil_h-2), round((-mil_h + 1)/2), round((-mil_h - 1)/2)], [0, 1, -1], [mil_h, mil_k, mil_l]]	# Slightly skewed, otherwise a very large cell is required
		single_divider = int( round( 3*(mil_h-2) + 2 ) )
	else:
		print('WARNING: miller index not implemented')
		exit(1)
elif mil_h == mil_k:			# 111 surface with 110 steps
	if (mil_h - mil_l) == 1:
		Nterrace = mil_h*2
		directions=[[1, -1, 0], [mil_l*2, mil_l*2, -mil_h*4], [mil_h, mil_k, mil_l]]
		single_divider = int( round( 3*mil_h - 1 ) )
		if single_divider % 2 == 0:
			single_divider = int( round(single_divider/2) )
	elif (mil_h - mil_l) == 2:
		Nterrace = mil_h
		directions=[[1, -1, 0], [round((mil_l + 1)/2), round((mil_l - 1)/2), -mil_h], [mil_h, mil_k, mil_l]]	# Slightly skewed, otherwise a very large cell is required
		single_divider = int( round( 3*(mil_h-1) ) )
	else:
		print('WARNING: miller index not implemented')
		exit(1)

# Required for multiplication of unit cell and removal of layers
divider = int( round( single_divider*(1 + layers // single_divider) ) )
FCCz = 1 + layers // single_divider

# Slab is cut from the bulk
slab = FaceCenteredCubic('Pt',
                              directions=directions,
                              miller=(None, None, (mil_h, mil_k, mil_l)),
                              latticeconstant=lattice,
                              size=(1, 1, FCCz),			# This multiplication is required if number of layers is >= divider,
                              pbc=True)					# otherwise the number of layers in the step is wrong

# Renumber systematically from top down.
orders = [(atom.index, round(atom.x, 3), round(atom.y, 3),
	-round(atom.z, 3), atom.index) for atom in slab]
orders.sort(key=itemgetter(3, 1, 2))
newslab = slab.copy()
for index, order in enumerate(orders):
	newslab[index].position = slab[order[0]].position.copy()
slab = newslab

def interlayer(slab, dz):
	"Corrects the outer 3 layers with the relaxed interlayer distance"

	if mil_h == 1 and mil_k == 1 and mil_l == 1:
		Nrow = supercellx*supercelly
		Nterrace_repeat = 1
	elif mil_k == mil_l:
		Nrow = supercelly
		Nterrace_repeat = supercellx
	elif mil_h == mil_k:
		Nrow = supercellx
		Nterrace_repeat = supercelly

	for i in range( 2, -1, -1 ):
		for j in range( Nterrace-1, -1, -1 ):
			idx = i*Nterrace + j - 1
			slab[idx].position[2] = slab[idx+Nterrace].position[2] + dz[i][j]

	return

# Here we correct the outer 3 layers with the relaxed interlayer distance
# Highest layer is always first, then within a layer the highest row is first
# All relaxations are based on a 1x1 super cell with 5 layers
if mil_h == 1 and mil_k == 1 and mil_l == 1:
	interlayer(slab,[[2.33615969537474],
			[2.28857266506526],
			[2.28778818208795]])
elif mil_h == 2 and mil_k == 1 and mil_l == 1:
	interlayer(slab,[[2.36432968046887,2.47062169750839,2.58223877378515],
			[2.382753627451,2.44507694770925,2.46712874389246],
			[2.41965547115621,2.44693538516596,2.42999456685702]])
elif mil_h == 5 and mil_k == 3 and mil_l == 3:
	interlayer(slab,[[2.36545824962164,2.46062551104144,2.46253835156006,2.56225185647453],
			[2.37893130197709,2.40920510888687,2.45164691527855,2.45591475691614],
			[2.40127096130736,2.41302494121551,2.45374385229251,2.42208622865898]])
elif mil_h == 3 and mil_k == 2 and mil_l == 2:
	interlayer(slab,[[2.35572422956875,2.44239002710608,2.45510336445311,2.44259658220022,2.54212038629775],
			[2.36479973548947,2.39344233933972,2.40682946980459,2.44224705255925,2.43304387701269],
			[2.38562856964255,2.39511513274094,2.41204807229824,2.43805574750581,2.40222877371326]])
elif mil_h == 7 and mil_k == 5 and mil_l == 5:
	interlayer(slab,[[2.34505484208797,2.43019193811589,2.44101762324542,2.43941394335883,2.42747043917829,2.52499803144303],
			[2.35142862562574,2.38179588556221,2.39187841936385,2.39580697834752,2.42383168341245,2.4165949686328],
			[2.3719201729736,2.38358267339244,2.39138433747394,2.39927767470594,2.4222987922369,2.38587764324528]])
elif mil_h == 4 and mil_k == 3 and mil_l == 3:
	interlayer(slab,[[2.33639395293959,2.41834758259237,2.42959555320907,2.42838437923218,2.42780456574947,2.41538215577957,2.5099858538219],
			[2.34062483713863,2.36994845714987,2.38092114828476,2.38382201194296,2.37983960087621,2.41283026935753,2.40274920585931],
			[2.35923177401331,2.37301380536016,2.38027073117343,2.37991213190783,2.38693616342181,2.41020487872846,2.37284395501425]])
elif mil_h == 9 and mil_k == 7 and mil_l == 7:
	interlayer(slab,[[2.32861283529512,2.40886429548785,2.42049456751182,2.41879355572371,2.41852628353562,2.41816675125983,2.40416282528838,2.49966589360593],
			[2.33200779504031,2.36088972318536,2.37089632762854,2.37467605624036,2.37057513102265,2.37022046953394,2.40342686189086,2.39200643186493],
			[2.3495052486443,2.36361677879752,2.37139539425991,2.37063351143102,2.36895083244267,2.37777955502236,2.40034037114573,2.36260689083348]])
elif mil_h == 2 and mil_k == 2 and mil_l == 1:
	interlayer(slab,[[2.60826688741779,2.67115062762637,2.69541795887316,2.75789449836208],
			[2.6011380860155,2.63989314525058,2.690521368743,2.65811766828198],
			[2.63498905145464,2.66224264203773,2.6632039816726,2.63201930888907]])
elif mil_h == 5 and mil_k == 5 and mil_l == 3:
	interlayer(slab,[[2.55104554600059,2.61321696687046,2.62479199386752,2.62875434262227,2.69323600948688],
			[2.53791224618212,2.57445932625144,2.58590224419014,2.63335712158168,2.58513565890589],
			[2.56554805607553,2.58636583069305,2.6047292914682,2.5946939576493,2.56718500339513]])
elif mil_h == 3 and mil_k == 3 and mil_l == 2:
	interlayer(slab,[[2.50640808775485,2.56655613468912,2.58564250774539,2.57878918178502,2.58685530867188,2.64110983665329],
			[2.49099975026026,2.53056058126831,2.53829876122644,2.54545314434424,2.58237498481103,2.53473692547469],
			[2.51876149817193,2.53817442011362,2.54637246518395,2.55618259306485,2.54522491527667,2.52033160013763]])
elif mil_h == 7 and mil_k == 7 and mil_l == 5:
	interlayer(slab,[[2.47310824910675,2.5338731912258,2.55097572305608,2.5523374483664,2.54943970338456,2.54742646893475,2.60339445388141],
			[2.45760411942751,2.49693528869372,2.50635413167963,2.51059336001378,2.50437770389347,2.54509972915038,2.49857065607301],
			[2.48541423679049,2.50406506951636,2.50870935251334,2.50857044725754,2.52008119348034,2.5089983569139,2.48658196265383]])
elif mil_h == 4 and mil_k == 4 and mil_l == 3:
	interlayer(slab,[[2.44689337021112,2.50890445045483,2.52668636132954,2.52434565548473,2.53153615098757,2.52027047639793,2.51734514266842,2.575013348548],
			[2.43244615757969,2.47195461062749,2.47945374728183,2.48584897198109,2.47946798362067,2.47448557804559,2.51786595311231,2.47156711480691],
			[2.46022486353961,2.47789075977741,2.48151744590346,2.47992655730596,2.4800280424079,2.49278871953773,2.48186673513751,2.46111153994796]])
elif mil_h == 9 and mil_k == 9 and mil_l == 7:
	interlayer(slab,[[2.42727994957764,2.48780408682867,2.5067900515555,2.50555682957485,2.5087937695235,2.50714174405929,2.49626392081144,2.49643864126281,2.55329471845184],
			[2.41198334025944,2.4512272082618,2.45985375892883,2.46419468029573,2.45991009607859,2.45509792815695,2.45252362089333,2.49776397782147,2.44978739128413],
			[2.43944769214493,2.45795440919207,2.46014512746596,2.45778169075708,2.45671392378963,2.45832155928258,2.47231462396765,2.46018218781556,2.44064813269961]])
else:
	print('WARNING: layer relaxation of this miller index is not implemented')
	exit(1)

slab = slab.repeat((supercellx, supercelly, 1))

# Renumber systematically from top down.
orders = [(atom.index, round(atom.x, 3), round(atom.y, 3),
	-round(atom.z, 3), atom.index) for atom in slab]
orders.sort(key=itemgetter(3, 1, 2))
newslab = slab.copy()
for index, order in enumerate(orders):
	newslab[index].position = slab[order[0]].position.copy()
slab = newslab

# Remove the additional layers
if layers % divider:
	for i in range(len(slab), (layers)*supercellx*Nterrace*supercelly, -1):
		del slab[-1]
	slab.wrap()
	dz = min(slab.positions[:,2])
	slab.translate((0., 0., -dz))
	slab.cell[2][2] -= dz

# Here we add vacuum
Cell = slab.get_cell()
if mil_h == 1 and mil_k == 1 and mil_l == 1:
	idx = 0
elif mil_k == mil_l:
	idx = supercellx*Nterrace*supercelly - supercelly
elif mil_h == mil_k:
	idx = supercellx*Nterrace*supercelly - supercellx
Cell[2][2] = slab[idx].position[2]
slab.set_cell(Cell)
add_vacuum(slab, vacuum)

# Translate all atoms in order to get the top step atom at the (0, 0, 0) position
slab.wrap()
idx = np.argmax(slab.positions[:,2]) 
slab.translate( -slab.positions[idx,:] )
slab.wrap(pbc=(1,1,0))

# Renumber systematically from top down, again.
orders = [(atom.index, round(atom.x, 3), round(atom.y, 3),
	-round(atom.z, 3), atom.index) for atom in slab]
orders.sort(key=itemgetter(3, 1, 2))
newslab = slab.copy()
for index, order in enumerate(orders):
	newslab[index].position = slab[order[0]].position.copy()
slab = newslab

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

def get_ZH(slab, X, Y):
	"""
	Used to estimate the height difference w.r.t. the step atom according to the slope of surface
	"""
	if mil_h == mil_l and mil_h == mil_k:	# 111 surface without steps
		ZH = 0.
	elif mil_k == mil_l:			# 111 surface with 100 steps
		midpoint = slab[supercellx*Nterrace*supercelly-supercelly].position
		midpoint_d = np.dot( midpoint, Cell_ )
		if X <= midpoint_d[0]:
			dzdx = -abs(midpoint[2]) / midpoint_d[0]
		else:
			dzdx = abs(midpoint[2]) / ( 1.0 - midpoint_d[0] )
		ZH = (X - midpoint_d[0]) * dzdx + midpoint[2]
	elif mil_h == mil_k:			# 111 surface with 110 steps
		midpoint = slab[supercellx*Nterrace*supercelly-supercellx].position
		midpoint_d = np.dot( midpoint, Cell_ )
		if Y <= midpoint_d[1]:
			dzdy = -abs(midpoint[2]) / midpoint_d[1]
		else:
			dzdy = abs(midpoint[2]) / ( 1.0 - midpoint_d[1] )
		ZH = (Y - midpoint_d[1]) * dzdy + midpoint[2]

	return ZH

def rand_conditions(Zmin, Zmax, rmin, rmax, slab):
	"""
	Sample randomly the initial conditions
	"""
	if mil_h == mil_l and mil_h == mil_k:
		X       = np.random.random_sample()
		Y       = np.random.random_sample()
	elif mil_k == mil_l:
		X       = np.random.random_sample() / supercellx
		Y       = np.random.random_sample()
	elif mil_h == mil_k:
		X       = np.random.random_sample()
		Y       = np.random.random_sample() / supercelly
	THETA   = np.random.random_sample() * 180.
	PHI     = np.random.random_sample() * 360.
	Z = (Zmax - Zmin) * np.random.random_sample() + Zmin + get_ZH(slab, X, Y)
	r = (rmax - rmin) * np.random.random_sample() + rmin

	return Z, r, X, Y, THETA, PHI

def position(Z, r, X, Y, THETA, PHI):
	"""
	Convert spherical coordinates to cartesian coordinates
	"""
	PHI     = np.radians(PHI)
	THETA   = np.radians(THETA)
	RotMatrix =	[[np.cos(PHI),                np.sin(PHI),               0.            ],
			[-np.cos(THETA)*np.sin(PHI),  np.cos(THETA)*np.cos(PHI), np.sin(THETA) ],
			[ np.sin(THETA)*np.sin(PHI), -np.sin(THETA)*np.cos(PHI), np.cos(THETA) ]]

	v_H = [ 0., 0., r ]
	v_H = np.dot( v_H, RotMatrix )

	H1 = np.dot([X, Y, Z], Cell)
	H1[2] = Z
	H2 = H1 + v_H

	return [H1, H2]

# Here we determine the configuration of H2 and add it to the system
H2 = molecule('H2')
Break_loop = False
while Break_loop == False:
	Z, r, X, Y, THETA, PHI = rand_conditions(Zmin, Zmax, rmin, rmax, slab)

	H2.positions = position(Z, r, X, Y, THETA, PHI)

	H_breakloop = [False, False]
	for i in range(2):
		H_d = np.dot( H2[i].position, Cell_ )
		if (H2[i].position[2] - get_ZH(slab, H_d[0], H_d[1])) > 0.3:
			H_breakloop[i] = True

	if H_breakloop[0] and H_breakloop[1]: Break_loop = True
add_adsorbate( slab, H2, height=Z, position=(H2[0].position[0], H2[0].position[1]) )

# These frequencies came from davide's SRP32-vdW calculations, should be good enough
Frequency =  [ [ 11.290474, 11.286502, 10.634782 ],
	        # omega X ,  omega Y  ,   omega Z  -- Layer 1 (meV)
               [ 13.744123, 13.746207, 13.595279 ],
	        # omega X ,  omega Y  ,   omega Z  -- Layer 2 (meV)
               [ 14.279348, 14.281448, 14.112210 ] ]
                # omega X ,  omega Y  ,   omega Z  -- Layer 3 (meV)

# 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)
Ts = 500.			# Surface temperature

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
# using a independent harmonic oscillator model
Potential = 0.
for i  in range( layers*supercellx*Nterrace*supercelly ):
	LayerNumber =  i // ( supercellx*Nterrace*supercelly )
	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.
			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 )) )

# Write a POSCAR file
#write('POSCAR', slab, format='vasp')
write('POSCAR', slab, format='vasp', direct='true')