from formulary import *
from util.aseutil import set_atom_color, get_pov_settings

"""
Create crystal structures using ase and render them with povray

Rotation angle:
To get the rotation angle, open the structure in the ase.visualize.view
and use "View->Rotation" to get the desired angles
"""

set_atom_color("Na", COLORSCHEME["fg-red"])
set_atom_color("Cl", COLORSCHEME["fg-blue"])

set_atom_color("Zn", COLORSCHEME["fg-blue"])
set_atom_color("S", COLORSCHEME["fg-yellow"])

from ase.lattice import compounds
from ase.build import cut, bulk
from ase import Atom, Atoms


def zincblende():
    zns = compounds.Zincblende(("Zn", "S"), latticeconstant=5.0, size=(1,1,1))
    zns_cell = cut(zns, b=(0,0,1), origo=(0,0,0), extend=1.1)
    return zns_cell

# NaCl cut
def nacl():
    nacl = compounds.NaCl(("Na", "Cl"), latticeconstant=5.0, size=(1,1,1))
    nacl_cell = cut(nacl, b=(0,0,1), origo=(0,0,0), extend=1.1)
    return nacl_cell

def wurtzite():
    compounds.L1_2
    wurtzite = bulk('SZn', 'wurtzite', a=3.129, c=5.017)
    wurtzite_cell = cut(wurtzite,
        a=[1, 0, 0],
        b=[-1, -1, 0],
        c=[0, 0, 1], extend=1.1)
    return wurtzite_cell


if __name__ == "__main__":
    export_atoms(nacl(), "cm_crystal_NaCl", size_formula_half_quadratic)
    export_atoms(wurtzite(), "cm_crystal_wurtzite", size_formula_half_quadratic, rotation="70x,20y,174z")
    export_atoms(zincblende(), "cm_crystal_zincblende", size_formula_half_quadratic, rotation="-155x,70y,24z")
    w = wurtzite()
    from ase.visualize import view
    view(w)