#!/usr/bin/env python
from ase import Atoms
from ase.build import surface, add_vacuum, bulk, fcc211, cut
from ase.io import read, write
from ase.constraints import FixAtoms, FixedLine
from ase.lattice.cubic import FaceCenteredCubic
from operator import itemgetter
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])
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 = 3*mil_h - 2
		if single_divider % 2 == 0:
			single_divider = 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 = 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 = 3*mil_h - 1
		if single_divider % 2 == 0:
			single_divider = 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]]
		single_divider = 3*(mil_h-1)
	else:
		print('WARNING: miller index not implemented')
		exit(1)

divider = single_divider*(1 + layers // single_divider)
FCCz = int( 1 + layers // single_divider )

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
slab = slab.repeat((supercellx, supercelly, 1))

add_vacuum(slab, vacuum)

# 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

# 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,:] )

Natoms = len(slab)
mask = Natoms * [True]
maskrelaxlayers = []
Nlayers_relax = layers - 2
for i in range( Nlayers_relax*supercellx*Nterrace*supercelly ):
	maskrelaxlayers.append( i )
	mask[i] = False

fixlayers = FixAtoms(mask=mask)
#line = FixedLine(maskrelaxlayers, (0,0,1))
#slab.set_constraint([fixlayers, line])
slab.set_constraint([fixlayers])

write('POSCAR', slab, format='vasp', direct='true')
write('POSCAR', slab, format='vasp')