# Import the basic libraries

# ASE system
import ase
from ase import Atom, Atoms
from ase import io
from ase.lattice.spacegroup import crystal
from __future__ import division

# Spacegroup/symmetry library
from pyspglib import spglib

# iPython utility function
from IPython.core.display import Image

# Import the tools from the qe-util package
from qeutil import RemoteQE
from qeutil.analyzers import plot_bands

# Access info
import host

qe=RemoteQE(label='SiC-elec',   # A name for the project
               kpts=[8,8,8],    # Dense k-space grid
               xc='pz',         # Exchange functional type in the name of the pseudopotentials
               pp_type='vbc',   # Variant of the pseudopotential
               pp_format='UPF', # Format of the pseudopotential files
               ecutwfc=70,      # Energy cut-off
               pseudo_dir='../../pspot',
               use_symmetry=True,
               procs=8)         # Use 8 cores for the calculation


print qe.directory

a=4.3366
cryst = crystal(['Si', 'C'],
                [(0, 0, 0), (0.25, 0.25, 0.25)],
                spacegroup=216,
                cellpar=[a, a, a, 90, 90, 90])

# Assign the calculator to our system
cryst.set_calculator(qe)

print "Stress tensor before and after change (Voigt notation, GPa):\n", cryst.get_stress()/ase.units.GPa

a=4.3341
cryst.set_cell([a,a,a],scale_atoms=True)

print cryst.get_stress()/ase.units.GPa

# Several special points in the FCC lattice
G1=[0,0,0]
G2=[1,1,1]
X=[0,1,0]
L=[1/2,1/2,1/2]
W=[1/2,1/4,3/4]
K=[sqrt(1/2),sqrt(1/2),0]
M=[1,1,0]


# Set the path along which we would like to plot the electronic structure
qpath=array([G1,X,G2,L])

# Name the special points for plotting
qpname=[u'Γ','X',u'Γ','L']


qe.set(nbnd=8)                 # Increase the number of bands to include excited states 
qe.set(qpath=qpath)            # Define the path for the band structure calculation
qe.set(points=100)             # Set number of points along the path
qe.set(nkdos=[15,15,15])       # Set the grid for reciprocal space integration of states 

# execute the calculation of band structure and electronic dos
qe.calculate(cryst,properties=['bands','edos'])

# Print the position and size of the band gap
print "Highest Ocupied Level  (HOL): %(HOL)8.2f eV \nLowest Unocupied Level (LUL): %(LUL)8.2f eV" % qe.results
print "Band gap: %.2f eV" % (qe.results['LUL'] - qe.results['HOL'])

figsize(9,5)
plot(qe.results['edos'][0],qe.results['edos'][1])
axvspan(qe.results['HOL'], qe.results['LUL'], alpha=0.15, color='k'); # Mark the band gap with the gray rectangle
title('Electronic DOS')
xlabel('Energy (eV)')
ylabel('DOS (a.u.)');

figsize(12,7)
ax1,ax2=plot_bands(qe.results, qpath=qpath, qpname=qpname, label='DFT calc.')
ax2.legend();
ax2.set_xlabel('Wave vector')
ax2.set_ylabel('Energy (eV)')
ax2.set_title('$\\beta$-SiC band structure')
ax1.set_xlabel('DOS (a.u.)');