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>";


### Kamiya added
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 values
EF0 = 7.5883  #eV extracted from VASP


#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 function
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

# Choose one of E0 and compare the result for 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()