from __future__ import division, unicode_literals, print_function

"""
Utility functions:
Currently creates an ASE slab from a pymatgen input 
structure
"""

import math

from pymatgen.core.surface import Slab
from pymatgen.core.surface import SlabGenerator
from pymatgen.io.ase import AseAtomsAdaptor

from ase.lattice.surface import surface

def get_ase_slab(pmg_struct, hkl=(1,1,1), min_thick=10, min_vac=10):
    """
    takes in the intial structure as pymatgen Structure object
    uses ase to generate the slab
    returns pymatgen Slab object
    
    Args:
	pmg_struct: pymatgen structure object
	hkl: hkl index of surface of slab to be created
	min_thick: minimum thickness of slab in Angstroms
	min_vac: minimum vacuum spacing 
    """
    ase_atoms = AseAtomsAdaptor().get_atoms(pmg_struct)
    pmg_slab_gen = SlabGenerator(pmg_struct, hkl, min_thick, min_vac)
    h = pmg_slab_gen._proj_height
    nlayers = int(math.ceil(pmg_slab_gen.min_slab_size / h))
    ase_slab = surface(ase_atoms, hkl, nlayers)
    ase_slab.center(vacuum=min_vac/2, axis=2)
    pmg_slab_structure = AseAtomsAdaptor().get_structure(ase_slab)
    return Slab(lattice=pmg_slab_structure.lattice,
                species=pmg_slab_structure.species_and_occu,
                coords=pmg_slab_structure.frac_coords,
                site_properties=pmg_slab_structure.site_properties,
                miller_index=hkl, oriented_unit_cell=pmg_slab_structure,
                shift=0., scale_factor=None, energy=None)

@@ -162,52 +162,176 @@
 ort 
-math%0A%0Afrom pymatgen.core.surface import Slab
+sys%0Aimport math%0Aimport numpy as np%0A%0Afrom pymatgen.core.sites import PeriodicSite%0Afrom pymatgen.core.structure import Structure%0Afrom pymatgen.core.lattice import Lattice
 %0Afro
@@ -360,16 +360,22 @@
 e import
+ Slab,
  SlabGen
@@ -1619,8 +1619,1747 @@
 y=None)%0A
+%0A%0Adef slab_from_file(hkl, filename):%0A    %22%22%22%0A    reads in structure from the file and returns slab object.%0A    useful for reading in 2d/substrate structures from file.%0A    Args:%0A         hkl: miller index of the slab in the input file.%0A         filename: structure file in any format %0A                   supported by pymatgen%0A    Returns:%0A         Slab object%0A    %22%22%22%0A    slab_input = Structure.from_file(filename)%0A    return Slab(slab_input.lattice,%0A                slab_input.species_and_occu,%0A                slab_input.frac_coords,%0A                hkl,%0A                Structure.from_sites(slab_input,to_unit_cell=True),%0A                shift=0,%0A                scale_factor=np.eye(3, dtype=np.int),%0A                site_properties=slab_input.site_properties)%0A%0A%0Adef add_vacuum_padding(slab, vacuum):%0A    %22%22%22%0A    add vacuum spacing to the given structure%0A    Args:%0A        slab: sructure/slab object to be padded %0A        vacuum: in angstroms%0A    Returns:%0A         Structure object%0A    %22%22%22%0A    min_c = np.min(%5Bc%5B2%5D for c in slab.frac_coords%5D)%0A    slab.translate_sites(list(range(len(slab))), %5B0, 0, -min_c%5D)%0A    a, b, c = slab.lattice.matrix%0A    z = %5Bc%5B2%5D for c in slab.cart_coords%5D%0A    zmax = np.max(z)%0A    zmin = np.min(z)%0A    thickness = zmax - zmin%0A    new_c = c / np.linalg.norm(c) * (thickness+vacuum)%0A    new_lattice = Lattice(np.array(%5Ba,b,new_c%5D))%0A    new_sites = %5B%5D%0A    for site in slab:%0A        new_sites.append(PeriodicSite(site.species_and_occu,%0A                                      site.coords,%0A                                      new_lattice,%0A                                      properties=site.properties,%0A                                      coords_are_cartesian=True))%0A    return Structure.from_sites(new_sites)%0A

add utility functions for reading slab structures from file and adding vacuum padding