#!/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 os import path
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])

vacuum = 15.

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.
	Atoms is an atoms object and is edited in place.
	"""
	
	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)

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