'''
Uses the bandstructure file generated by Octopus and makes its information
available to use

Information contained in bandstructure file
-------------------------------------------
bandstructure -> number of kpoints, number of bands, CBM, VBM, condcution band,
                 and valence band
'''

import numpy as np
import numpy.ma as ma
import matplotlib.pyplot as plt
from glob import glob

from octopuspy.eigenvalues import Eigenvalues

class Bandstructure():
    '''
    Class that holds and gives methods to the information of a bandstructure
    file

    Attributes:
      _name (string): name of the semiconductor used for saving output files
      _bandstructure_path (string): filepath with the addition of the bandstructure file
      _bandstructure (numpy array): with shape (kpoints, band_info)
      _efermi_path (string): filepath with the addition of the total-dos-efermi.dat file
      efermi (float): fermi energy
      filepath (string): filepath to the bandstructure and total-dos-efermi.dat files
      kpoints (zipped numpy arrays): (kx array, ky array, kz array) shapes (num_kpoints, )
      num_bands (int): number of bands from the bandstructure file
      num_kpoints (int): number of kpoints from the bandstructure file
      energies (numpy array): with shape (num kpoints, num bands)
      occupancies (numpy array): with shape (num kpoints, num bands)
      occ_band_num (int): The number of occupied bands
    '''

    def __init__(self, filepath, name=None):
        '''
        Args:
          name (string): name of the semiconductor used for saving output files
          filepath (string): the filepath to the bandstructure and total-dos-efermi.dat files
        '''

        self._name = name
        self.filepath = filepath
        self.occ_band_num = None
        self._bandstructure_path = self.filepath + 'bandstructure'
        self._bandstructure = np.array(np.loadtxt(self._bandstructure_path))
        self._efermi_path = self.filepath + 'total-dos-efermi.dat'
        self.efermi = np.loadtxt(self._efermi_path)[0,0]
        self.kpoints = zip(self._bandstructure[:,1], self._bandstructure[:,2],
                           self._bandstructure[:,3])
        self.num_bands = self._bandstructure.shape[1] - 4
        self.num_kpoints = self._bandstructure.shape[0]
        self.energies = None
        self.occupancies = None
        self.ev = Eigenvalues(self.filepath, self.num_kpoints, self.num_bands)

    def get_eigenvalues(self):
        '''
        Isolates the bands from the bandstructure numpy array, subtracts the fermi
        energy and creates the energies and occupancies numpy arrays

        Params:
          bandstructure (numpy array): with shape (kpoints, band_info)
        Returns:
          energies (numpy array): with shape (num kpoints, num bands)
          occupancies (numpy array): with shape (num kpoints, num bands)
        '''

        # numpy array with shape (num kpoints, num bands)
        energies = self._bandstructure[:, 4:]
        energies = energies - self.efermi
        self.energies = energies

        self.occupancies = self.ev.get_occupancies()

        return(self.energies, self.occupancies)

    def plot_bands(self):
        '''
        Creates and saves a figure of the bandstructure plots. The valence band
        is colored darkorange and the conduction band is colored yellow.The
        valence band max and conduction band minimum are labeled with a '*'

        Output:
          bandstructure_plot.png: bandstructure figure saved to the gen_vasp folder
        '''

        x_data = self._bandstructure[:,0]

        occupied_bands, unoccupied_bands = self._split_bands()
        valence_band, vb_max = self._get_valence_band(occupied_bands)
        conduction_band, cb_min= self._get_conduction_band(unoccupied_bands)
        vb_max_index = valence_band.argmax()
        cb_min_index = conduction_band.argmin()

        fig = plt.figure()
        ax = fig.add_axes([0.1,0.1,0.8,0.8])

        # create plot for all but last row (valence band)
        # plot valence band on its own in color darkorange
        for band in occupied_bands[:-1,:]:
            ax.plot(x_data, band, color='darkred')
        ax.plot(x_data, valence_band, color='darkorange')

        # create plot for all but first row (conduction band)
        # plot conduction band on its own in yellow
        for band in unoccupied_bands[1:,:]:
            ax.plot(x_data, band, color='indigo')
        ax.plot(x_data, conduction_band, color='yellow')

        ax.plot(x_data[vb_max_index], vb_max, '*')
        ax.plot(x_data[cb_min_index], cb_min, '*')
        ax.set_ylabel('E-Ef')
        ax.set_title('Bulk Bandstructure')
        ax.set_xticklabels(['K','Gamma','X','W','K','Gamma','L','U','W','L','K'])
        # ax.set_xticklabels(['K','Gamma','X','W'])
        # ax.set_xticks([0.00,0.148,0.289,0.359])
        ax.tick_params(axis='both',labelsize=12)

        plt.axhline(y=0)
        if self._name is None:
            print('When using plot_bands you must provide a name to Bandstructure')
        else:
            fig.savefig('../gen_vasp/' + self._name + '/bandstructure_plot.png')

    def simple_plot(self):
        '''
        Plots the bands without identifying the valence or conduction bands
        '''

        fig = plt.figure()
        ax = fig.add_axes([0.1,0.1,0.8,0.8])

        x_data = self._bandstructure[:,0]

        for band in self.energies.T:
            ax.plot(x_data, band, color='indigo')

        ax.set_ylabel('E-Ef')
        ax.set_title('Bulk Bandstructure')
        ax.set_xticks([0.00,0.118,0.255,0.391,0.726,1.000])
        ax.set_xticklabels(['K','Gamma','X','W','K','Gamma','L','U','W','L','K'])
        ax.set_xticks([0.00,0.148,0.289,0.359,0.408,0.557,0.679,0.765,0.814,0.914,1.000])
        ax.tick_params(axis='both',labelsize=12)

        plt.axhline(y=0)
        if self._name is None:
            print('When using simple_plot you must provide a name to Bandstructure')
        else:
            fig.savefig('../gen_vasp/' + self._name + '/bandstructure_plot.png')

    def _get_conduction_band(self, unoccupied_bands):
        '''
        Separates the conduction band into its own numpy array

        Args:
          unoccupied_bands (numpy array): shape (num_bands, num_kpoints)
        Returns:
          conduction_band (numpy array): shape (num_kpoints, )
          cb_min (float): The smallest value in the conduction_band numpy array
        '''

        conduction_band = unoccupied_bands[0, :]
        cb_min = conduction_band.min()
        return(conduction_band, cb_min)

    def _get_valence_band(self, occupied_bands):
        '''
        Separates the valence band into its own numpy array

        Args:
          occupied_bands (numpy array): shape (num_bands, num_kpoints)
        Returns:
          valence_band (numpy array): shape (num_kpoints, )
          vb_max (float): The largest value in the valence_band numpy array
        '''

        try:
            valence_band = occupied_bands[-1, :]
            vb_max = valence_band.max()
        except IndexError as err:
            raise ValueError('Error determining valence band')

        return(valence_band, vb_max)

    def _split_bands(self):
        '''
        Takes the energies numpy array shape (num_bands, num_kpoints) and
        creates two numpy arrays, occupied and unoccupied bands.

        Returns:
          occupied_bands (numpy array): shape (num_bands, num_kpoints)
          unoccupied_bands (numpy array): shape (num_bands, num_kpoints)
        '''

        self.occ_band_num = self.ev.get_num_occ_bands()

        occupied_bands = self.energies[:,:self.occ_band_num]
        unoccupied_bands = self.energies[:,self.occ_band_num:]

        # change to shape (num_bands, num_kpoints)
        occupied_bands = occupied_bands.T
        unoccupied_bands = unoccupied_bands.T

        return(occupied_bands, unoccupied_bands)
