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_structure(vasp_xml_path):


    vasprun = Vasprun(vasp_xml_path)
    try:
        # Get the band structure
        band_structure = vasprun.get_band_structure(line_mode=True)
        k_points = [k.frac_coords for k in band_structure.kpoints]
        energies = band_structure.bands
        ek = list(energies.keys())
        for key in ek:
            if key == Spin.up:
                energies[1] = energies.pop(key)
            if key == Spin.down:
                energies[-1] = energies.pop(key)

        return band_structure, {"k_points": k_points, "energies": energies}
    except:
        return None


def get_band_gap(vasprun_xml_path):
    vasprun = Vasprun(vasprun_xml_path)
    band_structure = vasprun.get_band_structure()

    # Check if the material is a metal (no band gap)
    if band_structure.is_metal():
        return "Metallic character, no band gap"

    # Get band gap information
    gap_info = band_structure.get_band_gap()
    return gap_info


def plot_bands_structure(band_structure):
    plotter = BSPlotter(band_structure)
    plot = plotter.get_plot()
    plt.show()
    return plot

# vaspxml_path="/Users/voicutomut/Documents/GitHub/BespokeMaterials/DB/Databases/JARVIS/JARVIS2D/OPT-Bandst/JVASP-13600/vasprun.xml"
# band_structure, numerics =get_band_structure(vaspxml_path)
# print(band_structure)
# print(numerics )
# plot_bands_structure(band_structure)
# print(get_band_gap(vaspxml_path))
# # plot
# # import numpy as np
# # k_points=numerics["k_points"]
# # energies= numerics["energies"]
# # # Creating the plot
# # distances = [0]
# # for i in range(1, len(k_points)):
# #     # Calculate the distance between successive k-points
# #     distance = np.linalg.norm(k_points[i] - k_points[i-1])
# #     distances.append(distances[-1] + distance)
# #
# # print("energies:", energies)
# # ene=[b for b in  energies]
# # for b in  ene:
# #     bands=energies[b]
# #     for  b_ in bands:
# #         plt.plot(distances, b_)
# #
# # plt.xlabel("K-point Path")
# # plt.ylabel("Energy (eV)")
# # plt.title("Band Structure")
# # plt.show()