#!/usr/bin/python

# electronicdos.py v0.5 5-16-2012 Jeff Doak jeff.w.doak@gmail.com
import numpy as np
from scipy.interpolate import UnivariateSpline
from scipy.integrate import quad
from scipy.optimize import fsolve
import sys, subprocess

BOLTZCONST = 8.617e-5 #eV/K

class ElectronicDOS:
    """
    Class to calculate equilibrium carrier concentrations, as well as
    equilibrium thermodynamic properties of an electronic density of states.

    Class constants of ElectronicDOS:
    - BOLTZCONST - Boltzmann's Constant (eV/K)

    Instance attributes of ElectronicDOS:
    - n_atoms - number of atoms of the unit cell from which the DOS was
        calculated
    - energy - numpy array of energies at which DOS is calculated (eV)
    - dos_tot - numpy array of density of spin up and spin down allowed electron
        states at each energy in the array energy (# states/eV/atom)
    - dos_spin - numpy array of the difference in density between spin up and
        spin down states (# states/eV/atom)
    - e_min - minimum energy in numpy array energy (eV)
    - e_max - maximum energy in numpy array energy (eV)
    - e_fermi - zero Kelvin fermi energy for the electronic DOS (eV)
    - step_size - energy difference between two consecutive points in the DOSCAR
        file (eV)
    - vbm - valence band maximum, set to e_fermi for metals (eV)
    - cbm - conduction band minimum, set to e_fermi for metals (eV)
    - band_gap -  band gap around the fermi energy, zero for metals (eV)
    - temp - numpy array of temperatures at which to calculate equilibrium
        electron chemical potentials, electron concentrations, and hole
        concentrations (K)
    - mu_e - numpy array of electron chemical potentials calculated at each
        temperature in temp (eV)
    - num_e - numpy array of equilibrium electron concentrations calculated at
        each temperature in temp (# e's/atom)
    - num_h - numpy array of equilibrium hole concentrations calculated at each
        temperature in temp (# h's/atom)
    - E_el - numpy array of electronic energy calculated at each temperature
        in temp (eV/atom)
    - S_el - numpy array of electronic entropy calculated at each temperature
        in temp (kB/atom)
    - F_el - numpy array of electronic free energy calculated at each
        temperature in temp (eV/atom)
    """

    def __init__(self,input_,format=None):
        if isinstance(input_,str):
            try:
                input_ = open(input_,'r')
            except IOError:
                print "Error reading input file."
                print "Program will now exit!"
                sys.exit(1)
        if isinstance(input_,file):
            if format == "ezvasp":
                self.read_ezvasp_dos(input_)
            else:
                self.read_doscar(input_)
                nelec = subprocess.Popen("grep NELECT OUTCAR",
                    shell=True,stdin=None,stdout=subprocess.PIPE).communicate()[0]
                self.nelec = int(float(nelec.split()[2]))
            self.get_bandgap()
        # Calculate finite temperature properties
        self.temp = np.linspace(0,2000,21)
        self.mu_e = np.zeros_like(self.temp)
        self.num_e = np.zeros_like(self.temp)
        self.num_h = np.zeros_like(self.temp)
        self.E_el = np.zeros_like(self.temp)
        self.S_el = np.zeros_like(self.temp)
        self.F_el = np.zeros_like(self.temp)
        # Calculate E_el_0
        self.E_el_0 = None
        tol = 1e-5
        for i in range(len(self.temp)):
            if i < tol:
                self.mu_e[i] = self.e_fermi
                self.E_el[i] = 0.0
                self.S_el[i] = 0.0
                self.num_e[i] = 0.0
                self.num_h[i] = 0.0
            elif i > 0.0:
                self.mu_e[i] = self.calc_mu_e(self.temp[i])
                if self.E_el_0 == None:
                    self.E_el_0 = self.calc_E_el(self.mu_e[i],self.temp[i])
                self.num_e[i] = self.n(self.mu_e[i],self.temp[i])
                self.num_h[i] = self.p(self.mu_e[i],self.temp[i])
                self.E_el[i] = (self.calc_E_el(self.mu_e[i],self.temp[i]))
                self.S_el[i] = self.calc_S_el(self.mu_e[i],self.temp[i])

        self.E_el[1:] = self.E_el[1:] - self.E_el_0
        self.F_el = self.E_el - self.temp*BOLTZCONST*self.S_el

    def read_doscar(self,input_):
        """
        Reads in a doscar file to grab the density of states as a function of
        energy. The argument input_ is assumed to be a file object.
        """
        self.n_atoms = int(input_.readline().split()[0])
        # Discard header information
        for i in range(4):
            input_.readline()
        # Read in Fermi Energy
        line = input_.readline().split()
        self.e_max = float(line[0])
        self.e_min = float(line[1])
        self.e_fermi = float(line[3])
        energy = []; dos_tot = []; dos_spin = []
        for line in input_:
            line = line.split()
            energy.append(float(line[0]))
            if len(line) == 3:
                dos_tot.append(float(line[1]))  # DOS includes spin up and down
                dos_spin.append(0.0)
            elif len(line) == 5:
                dos_tot.append(float(line[1])+float(line[2]))
                dos_spin.append(float(line[1])-float(line[2]))
        self.energy = np.array(energy)
        #self.dos_tot = np.array(dos_tot)/float(self.n_atoms)
        self.dos_tot = np.array(dos_tot)
        #self.dos_spin = np.array(dos_spin)/float(self.n_atoms)
        self.dos_spin = np.array(dos_spin)
        self.dos_spline = UnivariateSpline(self.energy,self.dos_tot)

    def read_ezvasp_dos(self,input_):
        """
        Reads an ezvasp-formatted dos.out file to get the electronic density of
        states. The argument input_ is assumned to be a file object.
        """
        nions = subprocess.Popen("grep NIONS OUTCAR",
                shell=True,stdin=None,stdout=subprocess.PIPE).communicate()[0]
        self.n_atoms = int(float(nions.split()[-1]))
        self.e_min = 0.0
        line = input_.readline().split()
        self.nelec = int(float(line[0]))
        self.step_size = float(line[1])
        self.scale = float(line[2])
        energy = []; dos_tot = []
        i = 0
        for line in input_:
            line = line.split()
            dos_tot.append(float(line[0]))
            energy.append(float(i)*self.step_size)
            i += 1
        self.energy = np.array(energy)
        self.dos_tot = np.array(dos_tot)
        self.dos_spin = np.zeros_like(self.dos_tot) # Change this for spin-polar
        self.dos_spline = UnivariateSpline(self.energy,self.dos_tot)
        self.e_max = self.energy[-1]
        # Find the 0 Kelvin 'Fermi Energy' using ATAT's method
        ne = 0.0
        for i in range(len(self.dos_tot)):
            ne += self.dos_tot[i]*self.step_size
            e_fermi = self.energy[i]
            if ne >= self.nelec:
                break
        self.e_fermi = e_fermi

    def get_bandgap(self):
        """
        Finds the band gap of a DOS around the fermi energy.
        """
        self.step_size = self.energy[1] - self.energy[0]
        i = 0
        not_found = True
        while not_found:
            if self.energy[i] < self.e_fermi and self.dos_tot[i] > 1e-3:
                bot = self.energy[i]
            elif self.energy[i] > self.e_fermi and self.dos_tot[i] > 1e-3:
                top = self.energy[i]
                not_found = False
            i += 1
        if top - bot < 2*self.step_size:
            self.vbm = self.cbm = self.e_fermi
            self.band_gap = 0.0
        else:
            self.vbm = bot; self.cbm = top
            self.band_gap = top - bot

    def shift_energy(self,new_ref):
        """
        Change the reference energy for all of the energy attributes.
        """
        self.energy = self.energy - new_ref
        self.e_min = self.e_min - new_ref
        self.e_max = self.e_max - new_ref
        self.e_fermi = self.e_fermi - new_ref
        self.vbm = self.vbm - new_ref
        self.cbm = self.cbm - new_ref
        self.mu_e = self.mu_e - new_ref

    #def sum_dos(self,weight,start,end,args=None):
    def sum_dos(self,weight,start,end,args=None):
        """
        Sums the density of states, dos, in the energy range [start,end], weighted
        by the function weight, which takes as inputs energy and args.
        """
        flag = False
        sum = 0.
        for i in range(len(self.energy)):
            if flag:
                sum += self.step_size*self.dos_tot[i]*weight(
                        self.energy[i],args)
                if self.energy[i] > end:
                    break
            elif self.energy[i] >= start:
                flag = True
        return sum

    #def integrate_dos(self,weight,start,end,args=None,threshold=0.1):
    def ium_dos(self,weight,start,end,args=None,threshold=0.1):
        """
        Takes numpy arrays containing the energy and dos and integrates them over
        the range [start,end] with the weighting function weight. Weight should take
        as an argument the integrated energy and a list of other arguements args.
        """
        def integrand(x,weight,args):
            return self.dos_spline(x)*weight(x,args)
        result = quad(
                integrand,start,end,args=(weight,args),full_output=1,limit=350)
        integral = result[0]
        error = result[1]
        #if error > integral*threshold:
        #    print "Numerical integration error is greater than"
        #    print str(threshold)+" of the integrated value."
        #    sys.exit(1)
        return integral

    def n(self,mu_e,T):
        """
        Calculate the intrinsic number of conduction electrons per atom at an
        electron chemical potential mu_e and temperature T.
        """
        def fermi(x,args):
            mu = args[0]; T = args[1]
            return 1./(np.exp((x-mu)/(BOLTZCONST*T))+1.)
        #n = self.integrate_dos(fermi,self.cbm,self.e_max,args=(mu_e,T))
        #n = self.sum_dos(fermi,self.cbm,self.e_max,args=(mu_e,T))
        n = self.sum_dos(fermi,mu_e,self.e_max,args=(mu_e,T))
        return n

    def p(self,mu_e,T):
        """
        Calculate the intrinsic number of valence holes per atom at an electron
        chemical potential of mu_e and temperature T.
        """
        def fermi(x,args):
            mu = args[0]; T = args[1]
            return 1./(np.exp((mu-x)/(BOLTZCONST*T))+1.)
        #p = self.integrate_dos(fermi,self.e_min,self.vbm,args=(mu_e,T))
        #p = self.sum_dos(fermi,self.e_min,self.vbm,args=(mu_e,T))
        p = self.sum_dos(fermi,self.e_min,mu_e,args=(mu_e,T))
        return p

    def charge_neut2(self,mu_e,args):
        def fermi(x,args):
            mu = args[0]; T = args[1]
            return 1./(np.exp((x-mu)/(BOLTZCONST*T))+1.)
        T = args
        n_sum = self.sum_dos(fermi,self.e_min,self.e_max,args=(mu_e,T))
        return self.nelec - n_sum

    def charge_neutrality(self,mu_e,args):
        """
        Condition for charge neutrality for intrinsic doping in a perfect
        semiconductor. This function should be overwritten for a more
        complicated case. 
        """
        T = args  # Args could also include atomic chemical potentials.
        return self.p(mu_e,T) - self.n(mu_e,T)

    def calc_mu_e(self,temp):
        """
        Calculate the electron chemical potential at temperature temp using the
        condition of charge neutrality.
        """
        #mu_e = fsolve(self.charge_neutrality,self.e_fermi,args=(temp))
        mu_e = fsolve(self.charge_neut2,self.e_fermi,args=(temp))
        return mu_e

    def calc_E_el(self,mu_e,T):
        """
        Calculate the electronic energy at a temperature T and electron chemical
        potential mu_e.
        """
        def energy(x,args):
            return x
        def fermi_energy(x,args):
            mu = args[0]; T = args[1]
            if x-mu < -30.0*BOLTZCONST*T:
                return x
            elif x-mu > 30.0*BOLTZCONST*T:
                return 0.0
            else:
                return x/(np.exp((x-mu)/(BOLTZCONST*T))+1.)
        #E = self.integrate_dos(fermi_energy,self.e_min,self.e_max,args=(mu_e,T))
        #E_0 = self.integrate_dos(
        #        fermi_energy,self.e_min,self.e_max,args=(mu_e,T))
        E = self.sum_dos(fermi_energy,self.e_min,self.e_max,args=(mu_e,T))
        #E_0 = self.sum_dos(energy,self.e_min,self.e_fermi,args=None)
        return E
    
    def calc_S_el(self,mu_e,T):
        """
        Calculate the electronic entropy at an electron chemical potential mu_e
        and temperature T.
        """
        def weight(x,args):
            mu = args[0]; T = args[1]
            x = (x - mu)/(BOLTZCONST*T)
            f = 1.0/(np.exp(x)+1)
            if f > 1e-5 and (1.0 - f) > 1e-5:
                return -f*np.log(f)-(1.-f)*np.log(1.-f)
            else:
                return 0.0
            #f = -np.log(np.exp(x)+1)/(np.exp(x)+1)
            #f += -np.log(np.exp(-x)+1)/(np.exp(-x)+1)
            #return f
        #S = self.integrate_dos(weight,self.e_min,self.e_max,args=(mu_e,T))
        S = self.sum_dos(weight,self.e_min,self.e_max,args=(mu_e,T))
        return S

def fermi_dirac_dist(x,args):
    """
    Calculates the Fermi-Dirac distribution for an energy x, temperature
    args[0], and electron chemical potential args[1].
    """
    T = args[0]; mu = args[1]
    return 1./(np.exp((x-mu)/(BOLTZCONST*T))+1.)

def test2(argv):
    doscar = ElectronicDOS(open(str(argv[0]),'r'))
    T = 500
    #n = doscar.integrate_dos(
    #        fermi_dirac_dist,doscar.cbm,doscar.e_max,args=(T,doscar.e_fermi))
    p = doscar.p(doscar.e_fermi,T)
    print p

def test3(argv):
    format = None
    if len(argv) > 1:
        format = str(argv[1])
    doscar = ElectronicDOS(open(str(argv[0]),'r'),format)
    print doscar.temp
    print doscar.num_e
    print doscar.num_h
    print doscar.E_el
    print doscar.S_el
    print doscar.F_el

def atat_test(argv):
    format = None
    if len(argv) > 1:
        format = str(argv[1])
    doscar = ElectronicDOS(open(str(argv[0]),'r'),format)
    print doscar.E_el_0
    for i in range(len(doscar.temp)):
        print doscar.temp[i],doscar.mu_e[i],doscar.E_el[i],doscar.S_el[i],doscar.F_el[i]


def test1(argv):
    import matplotlib.pyplot as plt
    doscar = ElectronicDOS(open(str(argv[0]),'r'))
    plt.plot(doscar.energy,doscar.dos_tot)
    plt.show()

def main(argv):
    import matplotlib.pyplot as plt
    doscar = open(str(argv[0]))
    e_fermi,energy,n_tot,n_spin = read_doscar(doscar)
    plt.plot(energy,n_tot)
    if len(argv) > 1:
        doscar2 = open(str(argv[1]))
        e_fermi2,energy2,n_tot2,n_spin2 = read_doscar(doscar2)
        plt.plot(energy2,n_tot2)
    plt.show()

if __name__ == "__main__":
    import sys
    #test3(sys.argv[1:])
    atat_test(sys.argv[1:])
