"""
Calculate the contributes of electrons for finite temperature energy and entropy
from DOSCAR of VASP

Developed by Tianping  Yin based on EF-T-metal.py,
modified by T. Kamiya in 2019
"""

import os
import sys
import numpy as np
from numpy import arange
from scipy import integrate
from scipy import optimize
from scipy.interpolate import interp1d
from matplotlib import pyplot as plt

# constants
e    = 1.60218e-19      # C
kB   = 1.380658e-23     # JK<sup>-1</sup>


Debug = 0


# Relative accuracy of the quad functin
eps = 1.0e-5
nmaxdiv = 150

# Temperature range
Tmin  =    0  # K
Tmax  = 1500
Tstep =   50
nT = int((Tmax - Tmin) / Tstep + 1)

#set EF in OUTCAR of VASP
EF0 = 7.5883  # eV


#read from DOSCAR
file = "DOSCAR"
E_raw, dos_raw = np.genfromtxt(file, skip_header=6, usecols=(0,1), unpack=True)
nDOS = len(E_raw)
Emin = E_raw[0]
Emax = E_raw[nDOS-1]
Estep = E_raw[1] - E_raw[0]


# define the DOS function
def f_dos(E):
    global E_raw, dos_raw
    global Emin, Emax, Estep, nDOS

    if E < Emin or Emax < E:
        print("Error in f_dos: Given E={} exceeds the DOS E range [{}, {}]".format(E, Emin, Emax))
        exit()

    fdos = interp1d(E_raw, dos_raw, kind = 'cubic')
    return fdos(E)

# Define the Fermi-Dirac distribution functions
def fe(E, T, EF):
    global e, kB

    if T == 0.0:
        if E < EF:
            return 1.0
        else:
            return 0.0
    return 1.0 / (np.exp((E - EF) * e / kB / T) + 1.0)

def fh(E, T, EF):
    return 1.0 - fe(E, T, EF)

# Define the function to be integrated
def dosfe(E, T, EF):
    return f_dos(E) * fe(E, T, EF)

# Calculate the number of electrons from E0 to E1 with the Fermi level EF at T
def Ne(T, EF, E0, E1):
    if T == 0.0:
        Ne_temp2 = integrate.quad(lambda E: f_dos(E), E0, EF, limit = nmaxdiv, epsrel = eps)
    else:
        Ne_temp2 = integrate.quad(lambda E: dosfe(E, T, EF), E0, E1, limit = nmaxdiv, epsrel = eps)
    if Debug:
        print("  Ne integ error: {}".format(Ne_temp2[1]))
    return Ne_temp2[0]

# Calculate the number of electrons from E0 to E1 (must be EF0) at 0 K,
def Ne0(EF0, E0, E1):
    return Ne(0.0, EF0, E0, E1)


# Calculate EF base on
# the charge neutrality (electron number) condition, Ne(T, EF, dE) - N = 0,
# using the Newton method with the initial value EF0
def CalEF(T, EF0, E0, E1, totNe):
    EF = optimize.newton(lambda EF: Ne(T, EF, E0, E1) - totNe, EF0)
    return EF

# Electron energy
def E_int(T, EF, E0, E1):
    Eint1 = integrate.quad(lambda E: dosfe(E, T, EF) * E, E0, E1, limit = nmaxdiv, epsrel = eps)
    if Debug:
        print(" E_int integ errors: {}".format(Eint1[1]))
    return Eint1[0]

# Mixing term in entropy without the factor of kB
def S(T, EF, E0, E1):
    if T == 0.0:
        return 0.0

    S1 = integrate.quad(lambda E: f_dos(E) * fe(E, T, EF) * np.log(fe(E, T, EF)), E0, E1, limit = nmaxdiv, epsrel = eps)
    S2 = integrate.quad(lambda E: f_dos(E) * fh(E, T, EF) * np.log(fh(E, T, EF)), E0, E1, limit = nmaxdiv, epsrel = eps)
    if Debug:
        print(" S integ errors: {}, {}".format(S1[1], S2[1]))
    return S1[0] + S2[0]

# Entropy * T in eV
def entropy(T, EF, E0, E1):
    return -T * kB / e * S(T, EF, E0, E1)


def main():
    global EF0

    print("quad() parameters")
    print("  releps=", eps)
    print("  nmaxdiv=", nmaxdiv)
    print("")
    print("Temperature range: {} - {}, {} K step".format(Tmin, Tmax, Tstep))
    print("")
    print("DOS configuration")
    print("  Read from [{}]".format(file))
    print("  E range: {} - {}, {} eV step".format(Emin, Emax, Estep))
    print("")
    print("EF at 0K: {} eV".format(EF0))
    Ne0K = Ne0(EF0, Emin, EF0)
    Eint0 = E_int(0.0, EF0, 0.0, EF0)
    print("  total Ne at 0 K: ", Ne0K)
    print("  total E at 0 K: ", Eint0, " eV")
    print("")

    xT = []
    yEF = []
    yEint = []
    yEentropy = []
    yEFa = []
    EFprev = EF0
    print ("%5s\t%8s\t%14s\t%14s\t%14s\t%14s" 
            % ("T(K)", "EF(eV)", "dE_ele(eV/u.c.)", "dS_ele*T(eV/u.c.) (approx)", "dF_ele(eV/u.c.)", 
               "Ne(0 K) (check)"))
    for i in range(nT):
        T = Tmin + i * Tstep
        kBTe = kB * T / e
        dE = 6 * kBTe

# Integration lowerlimit: Choose one of E0 and compare the result for debugging
# The first one is the most robust for bug, but very slow
# The third one is faster, but use after careful debugging
#        E0 = 0.0
#        E0 = Emin
        E0 = EFprev - dE
        E1 = EFprev + dE
        
        if E0 < Emin or Emax < E1:
            print("*** Integration range error: [{}, {}] exceeds DOS E range [{}, {}]".format(E0, E1, Emin, Emax))
            exit()

        Ne0K = Ne0(EF0, E0, EF0)
        Eint0 = E_int(0.0, EF0, E0, EF0)

        EF = CalEF(T, EFprev, E0, E1, Ne0K)
        NeCheck = Ne(T, EF, E0, E1)

        Eint = E_int(T, EF, E0, E1) - Eint0
        E_entropy = entropy(T, EF, E0, E1)
        F_electron = Eint + E_entropy
        
# The contribution of electrons to dS is roughly in the range [EF - kBT, EF + kBT] 
# The maximum contribution to S is at EF with fe = fh = 0.5
        E_entropy_approx = -kBTe * f_dos(EF) * (2.0 * kBTe) * 2.0 * 0.5 * np.log(0.5)

        xT.append(T)
        yEF.append(EF)
        yEint.append(Eint)
        yEentropy.append(E_entropy)
        yEFa.append(F_electron)

        EFprev = EF
        print ("%5.0f\t%12.6g\t%14.6g\t%14.6g(%14.6g)\t%14.6g\t%10.4g(%10.4g)" 
                % (T, EF, Eint, E_entropy, E_entropy_approx, F_electron, Ne0K, NeCheck))


#=============================
# Plot graphs
#=============================
    fig = plt.figure(figsize = (10, 10))

    ax1 = fig.add_subplot(2, 2, 1)
    ax2 = fig.add_subplot(2, 2, 2)
    ax3 = fig.add_subplot(2, 2, 4)
    ax1.plot(xT, yEint, label = 'E_int', color = 'navy')
    ax1.plot(xT, yEentropy, label = 'E_entropy', color = 'green')
    ax1.plot(xT, yEFa, label = 'Thermal_electron', color = 'red')
    ax1.set_xlabel("T (K)")
    ax1.set_ylabel("E (eV)")
    ax1.legend()
    ax2.plot(xT, yEF, label = 'EF', color = 'black')
    ax2.set_xlabel("T (K)")
    ax2.set_ylabel("EF (eV)")
    ax2.legend()
    ax3.plot(E_raw, dos_raw, label = 'DOS', color = 'black')
    ax3.plot([EF0, EF0], [min(dos_raw), max(dos_raw)], color = 'red', linestyle = 'dashed')
    ax3.set_xlabel("E (eV)")
    ax3.set_ylabel("DOS (states/u.c.)")
    ax3.legend()
# Rearange the graph axes so that they are not overlapped
    plt.tight_layout()

    plt.pause(0.1)


    print("Press ENTER to exit>>", end = '')
    input()


if __name__ == '__main__':
    main()