# Path: pyDFTutils/wannier90/band_plot.py
# def plot_band_weight(kslist,ekslist,wkslist=None,efermi=None,yrange=None,output=None,style='alpha',color='blue',axis=None,width=10,xticks=None):
#     if axis is None:
#         fig,a = plt.subplots()
#     else:
#         a=axis
# 
#     xmax=max(kslist[0])
#     if yrange is None:
#         yrange=(np.array(ekslist).flatten().min()-0.66,np.array(ekslist).flatten().max()+0.66)
# 
#     if wkslist is not None:
#         for i in range(len(kslist)):
#             x=kslist[i]
#             y=ekslist[i]
#             lwidths=np.array(wkslist[i])*width
#             #lwidths=np.ones(len(x))
#             points = np.array([x, y]).T.reshape(-1, 1, 2)
#             segments = np.concatenate([points[:-1], points[1:]], axis=1)
#             if style=='width':
#                 lc = LineCollection(segments, linewidths=lwidths,colors=color)
#             elif style=='alpha':
#                 lc = LineCollection(segments,linewidths=[4]*len(x), colors=[colorConverter.to_rgba(color,alpha=lwidth/(width+0.001)) for lwidth in lwidths])
# 
#             a.add_collection(lc)
#     if axis is None:
#         for ks , eks in zip(kslist,ekslist):
#             plt.plot(ks,eks,color='gray')
#         a.set_xlim(0,xmax)
#         a.set_ylim(yrange)
#         if xticks is not None:
#             plt.xticks(xticks[1],xticks[0])
#         if efermi is not None:
#             plt.axhline(linestyle='--',color='black')
#     return a


from pythtb import tb_model,w90
from ase.calculators.interface import Calculator,DFTCalculator
from ase.dft.dos import DOS
from ase.dft.kpoints import monkhorst_pack
from occupations import Occupations
from pyDFTutils.wannier90.band_plot import plot_band_weight
import numpy as np
import matplotlib.pyplot as plt
        """
        self.energy=(self._kweights*(self._occupations*self._eigenvals)).sum()

    def get_free_energy(self):
        pass
        #raise NotImplementedError

    def get_projection(self,orb,spin=0):
        """
        get the projection to nth orb.

        :param orb: the index of the orbital.
        :param spin: if spin polarized, 0 or 1

        :returns: eigenvecs[iband,ikpt]
        """
        if self._nspin==2:
            return self._eigenvecs[:,:,orb,spin]
        else:
            return self._eigenvecs[:,:,orb]

    def plot_projection(self,orb,spin=0,color='blue',axis=None):
        """
        plot the projection of the band to the basis
        """
        kslist=[list(range(len(self._kpts)))]*self._norb
        ekslist=self._eigenvals
        wkslist=np.abs(self.get_projection(orb,spin=spin))

        #fig,a = plt.subplots()

# Path: pyDFTutils/wannier90/band_plot.py
# def plot_band_weight(kslist,ekslist,wkslist=None,efermi=None,yrange=None,output=None,style='alpha',color='blue',axis=None,width=10,xticks=None):
#     if axis is None:
#         fig,a = plt.subplots()
#     else:
#         a=axis
# 
#     xmax=max(kslist[0])
#     if yrange is None:
#         yrange=(np.array(ekslist).flatten().min()-0.66,np.array(ekslist).flatten().max()+0.66)
# 
#     if wkslist is not None:
#         for i in range(len(kslist)):
#             x=kslist[i]
#             y=ekslist[i]
#             lwidths=np.array(wkslist[i])*width
#             #lwidths=np.ones(len(x))
#             points = np.array([x, y]).T.reshape(-1, 1, 2)
#             segments = np.concatenate([points[:-1], points[1:]], axis=1)
#             if style=='width':
#                 lc = LineCollection(segments, linewidths=lwidths,colors=color)
#             elif style=='alpha':
#                 lc = LineCollection(segments,linewidths=[4]*len(x), colors=[colorConverter.to_rgba(color,alpha=lwidth/(width+0.001)) for lwidth in lwidths])
# 
#             a.add_collection(lc)
#     if axis is None:
#         for ks , eks in zip(kslist,ekslist):
#             plt.plot(ks,eks,color='gray')
#         a.set_xlim(0,xmax)
#         a.set_ylim(yrange)
#         if xticks is not None:
#             plt.xticks(xticks[1],xticks[0])
#         if efermi is not None:
#             plt.axhline(linestyle='--',color='black')
#     return a


# Path: pyDFTutils/tightbinding/mypythTB.py
from pythtb import tb_model,w90
from ase.calculators.interface import Calculator,DFTCalculator
from ase.dft.dos import DOS
from ase.dft.kpoints import monkhorst_pack
from occupations import Occupations
from pyDFTutils.wannier90.band_plot import plot_band_weight
import numpy as np
import matplotlib.pyplot as plt

        """
        self.energy=(self._kweights*(self._occupations*self._eigenvals)).sum()

    def get_free_energy(self):
        pass
        #raise NotImplementedError

    def get_projection(self,orb,spin=0):
        """
        get the projection to nth orb.

        :param orb: the index of the orbital.
        :param spin: if spin polarized, 0 or 1

        :returns: eigenvecs[iband,ikpt]
        """
        if self._nspin==2:
            return self._eigenvecs[:,:,orb,spin]
        else:
            return self._eigenvecs[:,:,orb]

    def plot_projection(self,orb,spin=0,color='blue',axis=None):
        """
        plot the projection of the band to the basis
        """
        kslist=[list(range(len(self._kpts)))]*self._norb
        ekslist=self._eigenvals
        wkslist=np.abs(self.get_projection(orb,spin=spin))

        #fig,a = plt.subplots()

        return plot_band_weight(kslist,ekslist,wkslist=wkslist,efermi=None,yrange=None,output=None,style='alpha',color=color,axis=axis,width=10,xticks=None)