class Band():
    def write_data(self):
        if self.nspin==1:
            fout=file('band.data','w')
            for i in range(self.ndiv):
                fout.write("%12.8f "%self.band[0][i])
                for j in range(self.nband):
                    fout.write("%12.8f "%self.band[j+1][i])
                fout.write('\n')
            fout.close()
        elif self.nspin==2:
            foutup=file('bandup.data','w')
            foutdn=file('banddn.data','w')
            for i in range(self.ndiv):
                foutup.write("%12.8f "%self.band[0][i])
                foutdn.write("%12.8f "%self.band[0][i])
                for j in range(self.nband/2):
                    foutup.write("%12.8f "%self.band[j+1][i])
                foutup.write('\n')
                for j in range(self.nband/2,self.nband):
                    foutdn.write("%12.8f "%self.band[j+1][i])
                foutdn.write('\n')
            foutup.close()
            foutdn.close()

    def view_band(self,emin,emax,title,xlabel,ylabel):
        import numpy as np
        import matplotlib.pyplot as plt
        #plot setup
        plt.title(title)
        plt.xlabel(xlabel)
        plt.ylabel(ylabel)
        plt.xlim(self.kpos[0],self.kpos[-1])
        plt.ylim(emin,emax)
        plt.xticks(self.kpos,self.kpname)
        plt.grid(linestyle='-')
        plt.gca().xaxis.grid(True)
        plt.gca().yaxis.grid(False)
        #plot bands
        for i in range(self.nband):
            x=self.band[0]
            y=self.band[i+1]
            if np.min(y)emax:
                continue
            else:
                if self.nspin==1:
                    plt.plot(x,y,'r-')
                elif self.nspin==2:
                    if i                         plt.plot(x,y,'r-')
                    else:
                        plt.plot(x,y,'b--')
        plt.plot([self.kpos[0],self.kpos[-1]],[0.0,0.0],\
                 color='black',linestyle='--')
        plt.text(0.0,0.1,r'$E_F$=%6.3f eV'%self.fermi)
        plt.show()

class ElkBand(Band):
    def __init__(self):
        import sys
        import numpy as np
        try:
            fmain=file('elk.in','r')
            fpos=file('BANDLINES.OUT','r')
            fband=file('BAND.OUT','r')
            ffermi=file('EFERMI.OUT','r')
        except IOError:
            print 'Error in open nessary inputs. Check your input file'
            sys.exit()
        # initalize
        self.kpname=[] #names of kpoints
        self.kpos=[] #positions of special kpoints in G space
        self.nspin=1 #number of spins
        self.nband=0 #number of band
        self.band=[] #banddata
        # get self.nspin self.kpname self.nkp(number of special kpoints)
        # self.ndiv
        while True:
            tmp=fmain.readline()
            if len(tmp)==0:
                break
            if tmp.strip()=='spinpol':
                if fmain.readline().strip()=='.true.':
                    self.nspin=2
            if tmp.strip()=='plot1d':
                self.nkp,self.ndiv=fmain.readline().split()[0:2]
                self.nkp=int(self.nkp)
                self.ndiv=int(self.ndiv)
                for i in range(self.nkp):
                    self.kpname.append('$'+str(fmain.readline().split()[3])+'$')
        # get self.kpos
        while True:
            tmp=fpos.readline()
            if len(tmp)==0:
                break
            self.kpos.append(float(tmp.split()[0]))
            fpos.readline()
            fpos.readline()
        # get self.fermi
        self.fermi=27.2114*float(ffermi.readline().strip())
        # get self.nband self.band
        while True:
            tmp=fband.readline()
            if len(tmp)==0:
                break
            if len(tmp.strip())==0 and len(tmp)!=0:
                self.nband=self.nband+1
        rawband=np.loadtxt('BAND.OUT')
        kp=rawband[0:self.ndiv,0]
        self.band.append(kp)
        for i in range(self.nband):
            eig=27.2114*rawband[i*self.ndiv:(i+1)*self.ndiv,1]
            self.band.append(eig)
        #close files
        fmain.close()
        fpos.close()
        fband.close()
        ffermi.close()

    def write_data(self):
        Band.write_data(self)

    def view_band(self,emin=-10.0,emax=15.0,title='Band Plot',\
                 xlabel='K-Path',ylabel='Energy / eV'):
        Band.view_band(self,emin,emax,title,xlabel,ylabel)