import pymatgen.core as mg
from pymatgen.io.vasp.outputs import BSVasprun, Vasprun
from pymatgen.electronic_structure.core import Spin
from pymatgen.electronic_structure.plotter import BSPlotter, BSDOSPlotter, DosPlotter

import matplotlib.pyplot as plt
%matplotlib inline

"""
Explanation: Band Structure of NiO
Germain Salvato Vallverdu germain.vallverdu@univ-pau.fr
Short examples to show how to extract or plot the band structure from a VASP calculation using pymagen.
Written using:
- pymatgen==2018.3.13
First, load used packages :
End of explanation
"""


ls

"""
Explanation: I work in a directory that contains the calculations of the NiO band structure. Look at this example to learn how to obtain a band structure from VASP calculations.
End of explanation
"""


run = BSVasprun("vasprun.xml", parse_projected_eigen=True)

bs = run.get_band_structure("KPOINTS")

"""
Explanation: 1) Read vasprun.xml and extract the band structure
End of explanation
"""


print("number of bands", bs.nb_bands)

print("number of kpoints", len(bs.kpoints))

bs.is_metal()

bs.is_spin_polarized

"""
Explanation: You can obtain some information about the band structure :
End of explanation
"""


bs.bands

bs.bands[Spin.up].shape

"""
Explanation: 2) Extract a given band or eigenvalues
The bands attribute of the BaandStructure object is a dictionnary of arrays that contains all the bands. The shape is the following :
{Spin.up: np.array((nb_bands, nb_kpoints)), Spin.down: np.array((nb_bands, nb_kpoints))}
End of explanation
"""


bs.bands[Spin.down][9, :]

"""
Explanation: The 9th bands of spin down is extracted by :
End of explanation
"""


n = 0
for kpoints, e in zip(bs.kpoints, bs.bands[Spin.down][9, :]):
    n += 1
    if n == 11:
        print("...")
    if 10 < n < 190:
        continue

    print("kx = %5.3f  ky = %5.3f  kz = %5.3f   eps(k) = %8.4f" % (tuple(kpoints.frac_coords) +  (e,)))


"""
Explanation: In order to print a band and the corresponding k-points :
End of explanation
"""


bsplot = BSPlotter(bs)

# get the plot
bsplot.get_plot(ylim=(-20, 10), zero_to_efermi=True)
print(bs.efermi)

# add some features
ax = plt.gca()
ax.set_title("NiO Band Structure", fontsize=20)
xlim = ax.get_xlim()
ax.hlines(0, xlim[0], xlim[1], linestyles="dashed", color="black")

# add legend
ax.plot((), (), "b-", label="spin up")
ax.plot((), (), "r--", label="spin down")
ax.legend(fontsize=16, loc="upper left")

"""
Explanation: 3) Plot a simple band structure
End of explanation
"""


data = bsplot.bs_plot_data()

data.keys()

"""
Explanation: 3.1) Extract data from the plot
You can get data from the plot and in particular the (x, y) coordinates of each band.
End of explanation
"""


ibands = 9  # band number from 0 --> number of bands
spin = str(Spin.up)

for xpath, epath in zip(data["distances"], data["energy"]):
    print(20 * "-")
    for x, bands in zip(xpath, epath[spin][ibands]):
        print("%8.4f %8.4f" % (x, bands))

"""
Explanation: For example, here, you print the abscissa and the energy of the 9th band. Keep in mind that here the data are the ones used to do the plot. Thus the zero to fermi translation is already done according to the BSPlotter class.
End of explanation
"""


ibands = 9  # band number from 0 --> number of bands
spin = str(Spin.up)

for xpath, epath in zip(data["distances"], data["energy"]):
    plt.plot(xpath, epath[spin][ibands])

"""
Explanation: Same as above but in a plot.
End of explanation
"""


ibands = 1  # band number from 0 --> number of bands
spin = str(Spin.up)

x = list()
y = list()
for xpath, epath in zip(data["distances"], data["energy"]):
    x += xpath
    y += epath[spin][ibands]

plt.plot(x, y)

"""
Explanation: The same again but merging the slices of the band.
End of explanation
"""


dosrun = Vasprun("../DOS_SMEAR/vasprun.xml", parse_dos=True)
dos = dosrun.complete_dos
print(dosrun.efermi)
print(dos.efermi)

dosplot = DosPlotter(sigma=0.1)
dosplot.add_dos("Total DOS", dos)
dosplot.add_dos_dict(dos.get_element_dos())
plt = dosplot.get_plot()
plt.grid()

"""
Explanation: 4) Plot DOS
Read the DOS from another calculations.
End of explanation
"""


bs = run.get_band_structure("KPOINTS", efermi=dos.efermi)

bsdosplot = BSDOSPlotter(
    bs_projection="elements", 
    dos_projection="elements", 
    vb_energy_range=22,
    egrid_interval=2.5
)
plt = bsdosplot.get_plot(bs, dos=dos)

bsdosplot = BSDOSPlotter(
    bs_projection="elements", 
    dos_projection="elements", 
    vb_energy_range=11,
    cb_energy_range=10,
    egrid_interval=2
)
plt = bsdosplot.get_plot(bs, dos=dos)

"""
Explanation: 5) Plot Bands and DOS
End of explanation
"""


