"""
有限温度における電子のエネルギーとエントロピーへの寄与をVASPのDOSCARファイルから計算します。

このスクリプトは、VASPのDOSCARファイルから読み込んだ電子状態密度（DOS）データを使用して、
有限温度における電子の内部エネルギー、エントロピー、およびヘルムホルツ自由エネルギーを計算します。
フェルミ・ディラック分布を考慮し、温度に応じたフェルミ準位の変動もニュートン法を用いて決定します。
最終的に、温度に対するこれらの物理量の変化をプロットして可視化します。

:doc:`Electron_Helmholtz_energy_usage`
"""

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"


def f_dos(E):
    """
    与えられたエネルギーにおける電子の状態密度（DOS）を計算します。

    DOSCARファイルから読み込んだ離散的なDOSデータに対して、3次スプライン補間を用いて
    任意のエネルギーにおけるDOS値を求めます。

    :param E: エネルギー (float, eV)
    :returns: 指定されたエネルギーにおける状態密度 (float, states/eV/u.c.)
    """
    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)

def fe(E, T, EF):
    """
    フェルミ・ディラック分布関数を計算します。

    与えられたエネルギー、温度、フェルミ準位に基づいて、電子がそのエネルギー準位を占有する確率を計算します。
    絶対零度の場合には、E < EF で 1.0、それ以外で 0.0 を返します。

    :param E: エネルギー (float, eV)
    :param T: 温度 (float, K)
    :param EF: フェルミ準位 (float, eV)
    :returns: フェルミ・ディラック分布関数の値 (float)
    """
    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):
    """
    電子が存在しない確率 (1 - fe) を計算します。

    フェルミ・ディラック分布関数 `fe` を用いて、電子がそのエネルギー準位に存在しない確率を計算します。

    :param E: エネルギー (float, eV)
    :param T: 温度 (float, K)
    :param EF: フェルミ準位 (float, eV)
    :returns: 1 - fe(E, T, EF) の値 (float)
    """
    return 1.0 - fe(E, T, EF)

def dosfe(E, T, EF):
    """
    状態密度とフェルミ・ディラック分布関数の積を計算します。

    この関数は、積分計算で使用される被積分関数として定義されます。

    :param E: エネルギー (float, eV)
    :param T: 温度 (float, K)
    :param EF: フェルミ準位 (float, eV)
    :returns: f_dos(E) * fe(E, T, EF) の値 (float)
    """
    return f_dos(E) * fe(E, T, EF)

def Ne(T, EF, E0, E1):
    """
    指定されたエネルギー範囲における電子の総数を計算します。

    絶対零度ではフェルミ準位までのDOSを積分し、有限温度ではDOSとフェルミ・ディラック分布関数の積を積分します。

    :param T: 温度 (float, K)
    :param EF: フェルミ準位 (float, eV)
    :param E0: 積分下限エネルギー (float, eV)
    :param E1: 積分上限エネルギー (float, eV)
    :returns: 指定された範囲における電子の総数 (float)
    """
    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]

def Ne0(EF0, E0, E1):
    """
    絶対零度における電子の総数を計算します。

    `Ne` 関数を呼び出し、温度を0 Kに設定して絶対零度での電子数を計算します。

    :param EF0: 0 Kにおけるフェルミ準位 (float, eV)
    :param E0: 積分下限エネルギー (float, eV)
    :param E1: 積分上限エネルギー (float, eV)
    :returns: 絶対零度における電子の総数 (float)
    """
    return Ne(0.0, EF0, E0, E1)


def CalEF(T, EF0, E0, E1, totNe):
    """
    温度Tにおけるフェルミ準位EFを、電子数保持の条件から計算します。

    指定された温度Tにおいて、電子の総数が `totNe` となるようにフェルミ準位を調整します。
    ニュートン法 (`scipy.optimize.newton`) を用いて、`Ne(T, EF, E0, E1) - totNe = 0` を満たすEFを求めます。

    :param T: 温度 (float, K)
    :param EF0: 0 Kにおける初期フェルミ準位 (float, eV)
    :param E0: 積分下限エネルギー (float, eV)
    :param E1: 積分上限エネルギー (float, eV)
    :param totNe: 保持されるべき電子の総数 (float)
    :returns: 温度Tにおけるフェルミ準位 (float, eV)
    """
    EF = optimize.newton(lambda EF: Ne(T, EF, E0, E1) - totNe, EF0)
    return EF

def E_int(T, EF, E0, E1):
    """
    指定されたエネルギー範囲における電子の内部エネルギーを計算します。

    `f_dos(E) * fe(E, T, EF) * E` を指定されたエネルギー範囲で積分することにより、
    電子の内部エネルギーを求めます。

    :param T: 温度 (float, K)
    :param EF: フェルミ準位 (float, eV)
    :param E0: 積分下限エネルギー (float, eV)
    :param E1: 積分上限エネルギー (float, eV)
    :returns: 電子の内部エネルギー (float, eV/u.c.)
    """
    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]

def S(T, EF, E0, E1):
    """
    電子のエントロピーにおける混合項を、kBの係数を除いて計算します。

    フェルミ・ディラック分布関数 `fe` および `fh` を用いて、エントロピーの積分項を計算します。
    絶対零度ではエントロピーは0と見なされます。

    :param T: 温度 (float, K)
    :param EF: フェルミ準位 (float, eV)
    :param E0: 積分下限エネルギー (float, eV)
    :param E1: 積分上限エネルギー (float, eV)
    :returns: エントロピーの混合項 (float)
    """
    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]

def entropy(T, EF, E0, E1):
    """
    電子のエントロピーをエネルギー単位 (eV) で計算します。

    `S` 関数で計算された混合項に `-T * kB / e` の係数を乗じて、
    エントロピーを温度と電気素量でスケーリングしたエネルギー単位で返します。

    :param T: 温度 (float, K)
    :param EF: フェルミ準位 (float, eV)
    :param E0: 積分下限エネルギー (float, eV)
    :param E1: 積分上限エネルギー (float, eV)
    :returns: 電子のエントロピー (float, eV/u.c.)
    """
    return -T * kB / e * S(T, EF, E0, E1)


def main():
    """
    スクリプトのメイン実行関数です。

    VASPのDOSCARファイルから電子の状態密度を読み込み、指定された温度範囲で
    電子の内部エネルギー、エントロピー、およびヘルムホルツ自由エネルギーを計算します。
    結果はコンソールに出力され、matplotlibを用いてグラフ化されます。
    """
    global EF0

    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]


    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()