import pylab
import time
import math
import numpy
import pickle

from ase.io import *
import PROCAR

atoms=read("POSCAR")
#Read in PROCAR file
start_time = time.clock()
my_data=PROCAR.data("./PROCAR",1,False)

#Define energy shift
#Fermi energy
energy_shift = -3.9737
emin=-2.
emax=2.
width=0.05
npts=10001
offset=3
#energy_shift = 0.

pylab.figure(figsize=(17./2.54,20./2.54))
pylab.axes([0.15,0.15,0.6,0.80])
dos=my_data.get_dos(emin,emax,npts,width,energy_shift,False,"Gaussian",[0,])
shift=0.
atom_list=[61,60,59,58,57,56,1,2,4,5,3,0]
for iatom in range(len(atom_list)):
    shift=iatom*offset
    dpdos=my_data.get_partial_dos(emin,emax,npts,width,energy_shift,"Gaussian",[0,],[atom_list[iatom],],orbitals="d")
    pylab.plot(dpdos[:,0],dpdos[:,1]+shift,label=r"%s%d z = %.1f $\AA$"%(atoms[atom_list[iatom]].symbol,atom_list[iatom],atoms[atom_list[iatom]].position[2]),zorder=iatom)
shift=len(atom_list)*offset
pylab.plot(dos[:,0],dos[:,1]/10+shift,'black',label=r"DOS/10",zorder=2)

#pylab.legend(("Total","Mo s","Mo p", "Mo d","Sr s","Sr p", "Sr d","O s","O p"),loc='best',fontsize=8,ncol=2)
pylab.axvline(0.,color='k')
pylab.legend(loc=[1.05,0.5],fontsize=9)
pylab.xlim([emin,emax])
pylab.xlabel(r"$\mathsf{E-E_{Fermi}\,(eV)}$",fontsize=10)
pylab.ylabel("DOS",fontsize=10)
pylab.xticks(fontsize=9)
pylab.yticks(fontsize=9)
pylab.minorticks_on()
pylab.savefig("STO6SNO6_SCAN_Layer_by_Layer_PDOS_Fermi.png",dpi=600)

end_time = time.clock()
print(end_time-start_time)
