from pymatgen.io.vasp.outputs import Vasprun
from pymatgen.electronic_structure.plotter import BSPlotter
from pymatgen.electronic_structure.core import Spin
import matplotlib.pyplot as plt




def get_band_density_of_states(vasp_xml_path):
    vasprun = Vasprun(vasp_xml_path)

    # Extract total and projected DOS
    total_dos = vasprun.complete_dos

    # Access the energies and DOS values
    energies = total_dos.energies - total_dos.efermi  # Aligning energies with Fermi level
    total_dos_values = total_dos.densities

    return {"energies":energies,"total_dos_values":total_dos_values }

vasp_xml_path="/Users/voicutomut/Documents/GitHub/BespokeMaterials/DB/Databases/JARVIS/JARVIS2D/OPT-Bandst/JVASP-655/vasprun.xml"
b_dft=get_band_density_of_states(vasp_xml_path)
energies=b_dft["energies"]
total_dos_values=b_dft["total_dos_values"]
# Example: Plotting the total DOS
plt.plot(energies, total_dos_values[Spin.up], label='Spin Up')  # Spin-up DOS
if "Spin.down" in total_dos_values:
    plt.plot(energies, total_dos_values[Spin.down], label='Spin Down')  # Spin-down DOS

plt.xlabel('Energy (eV)')
plt.ylabel('DOS')
plt.title('Total Density of States')
plt.legend()
plt.show()