import sys
import time
from math import exp, sqrt

from scipy import integrate


"""
金属の有限温度 T における電子密度 N の計算。
数値積分の精度と計算時間を確認する。
"""


# 定数
e = 1.60218e-19       # C
kB = 1.380658e-23     # J K^-1


# Temperature
T = 300.0  # K

# Fermi energy
EF = 5.0  # eV

# EF近傍の nrange * kBT の範囲で積分
nrange = 6.0

# quad積分の相対積分精度と最大細分区間数
eps = 1.0e-8
maxorder = 100

# 計算時間を計測する繰り返し数
ncycle = 300


# コマンドライン引数の処理
if len(sys.argv) >= 2:
    T = float(sys.argv[1])

if len(sys.argv) >= 3:
    EF = float(sys.argv[2])


# eV単位のエネルギー差を E/(kBT) に変換する係数
ekBT = e / (kB * T)

# EF近傍の積分範囲
dE = nrange * kB * T / e


# 状態密度関数
def De(E):
    return sqrt(E)


# Fermi–Dirac分布関数
def fe(E, T, EF):
    if T == 0.0:
        if E < EF:
            return 1.0
        return 0.0

    return 1.0 / (exp((E - EF) * ekBT) + 1.0)


# 被積分関数
def Defe(E, T, EF):
    return De(E) * fe(E, T, EF)


def quad_integrate(func, Emin, Emax):
    """
    scipy.integrate.quadを用いて数値積分する。

    Returns
    -------
    value : float
        積分値
    error : float
        推定絶対誤差
    """
    return integrate.quad(
        func,
        Emin,
        Emax,
        epsrel=eps,
        limit=maxorder,
    )


def timed_quad(func, Emin, Emax):
    """
    quad積分をncycle回実行し、最後の積分結果と総計算時間を返す。
    """
    t0 = time.perf_counter()

    for _ in range(ncycle):
        ret = quad_integrate(func, Emin, Emax)

    elapsed_time = time.perf_counter() - t0

    return ret, elapsed_time


def main():
    print(f"T = {T} K   EF = {EF} eV")
    print()

    # De(E) = sqrt(E) の積分精度を確認
    Ianaly = 2.0 / 3.0 * EF**1.5

    print(
        "Analytical integration for De(E) "
        f"from 0.0 to {EF} = {Ianaly}"
    )

    ret = quad_integrate(De, 0.0, EF)
    print(f"quad: value = {ret[0]}, estimated error = {ret[1]}")
    print()

    integrand = lambda E: Defe(E, T, EF)

    # (1) 全積分範囲
    Emin, Emax = 0.0, EF + dE

    print(f"(1) Integration range: {Emin} - {Emax} eV")

    ret, elapsed_time = timed_quad(integrand, Emin, Emax)

    print(
        f"quad: value = {ret[0]}, "
        f"estimated error = {ret[1]}, "
        f"time = {elapsed_time} s"
    )
    print()

    # (2) Fermi準位より十分低い範囲
    Emin, Emax = 0.0, EF - dE

    print(f"(2) Integration range: {Emin} - {Emax} eV")

    Ianaly = 2.0 / 3.0 * (Emax**1.5 - Emin**1.5)

    print(
        "Analytical integration for De(E) "
        f"from {Emin} to {Emax} = {Ianaly}"
    )

    ret, elapsed_time = timed_quad(integrand, Emin, Emax)

    print(
        f"quad: value = {ret[0]}, "
        f"estimated error = {ret[1]}, "
        f"time = {elapsed_time} s"
    )
    print()

    # (3) Fermi準位近傍
    Emin, Emax = EF - dE, EF + dE

    print(f"(3) Integration range: {Emin} - {Emax} eV")

    ret, elapsed_time = timed_quad(integrand, Emin, Emax)

    print(
        f"quad: value = {ret[0]}, "
        f"estimated error = {ret[1]}, "
        f"time = {elapsed_time} s"
    )
    print()


if __name__ == "__main__":
    main()