# coding=utf-8
import os
import sys
from math import exp, sqrt
import numpy as np
from numpy import arange
from scipy import integrate         # 数値積分関数 integrateを読み込む
from scipy import optimize          # newton関数はscipy.optimizeモジュールに入っている
from scipy.interpolate import interp1d
from matplotlib import pyplot as plt


"""
金属のEFの温度依存性を数値積分とNewton法で計算。
近似式と比較。
Calculate EF(T) for metal using numerical integration and the Newton method,
and compare with the approximation formula
"""


# constants
pi   = 3.14159265358979323846
h    = 6.6260755e-34    # Js";
hbar = 1.05459e-34      # "Js";
c    = 2.99792458e8     # m/s";
e    = 1.60218e-19      # C";
kB   = 1.380658e-23     # JK<sup>-1</sup>";
me   = 9.1093897e-31    # kg";


### Kamiya added
Debug = 0


# Temperature range
Tmin  = 300  # K
Tmax  = 310
Tstep =  10
nT = int((Tmax - Tmin) / Tstep + 1)


#read from DOSCAR
file = "TDOS.dat"
E_raw, dos_raw = np.genfromtxt(file, skip_header=1, 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]

EF0 = 0.4
Egth = 0.01

Defects = [
            [+1.0, 0.1, 3.0],
            [-1.0, 0.2, 3.0]
          ]

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

#二分法パラメータ
#初期値範囲
dEbisec = 0.1
#EPS
epsbisec   = 1.0e-5
#最大繰り返し数
nmaxiter = 200
#繰り返し中に途中経過を出力するサイクル数
iprintiterval = 1

# integration parameters
# Integration range w.r.t. kBT
Einteg0 = -2.0
Einteg1 =  2.0
nrange  =  6.0



def FindBandEdges(E, DOS, EF0, Egth):
    EV = 1.0e30
    EC = 1.0e30
    i0 = -1
    for i in range(len(E)):
        if E[i] >= EF0:
            i0 = i
            break

    for i in range(i0, 0, -1):
        if DOS[i] > Egth:
            EV = E[i+1]
            break

    for i in range(i0, len(E), +1):
        if DOS[i] > Egth:
            EC = E[i-1]
            break

    return EV, EC


def rieman(x0, dx, y, xmin, xmax):
    S = 0.0
    for i in range(len(y)):
        x = x0 + i * dx
        if x < xmin or xmax < x:
            continue
        S += y[i]
    return S * dx

# define the DOS function
def 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)

# Fermi-Dirac distribution 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 / (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 DOS(E) * fe(E, T, EF)

def DOSfh(E, T, EF):
    return DOS(E) * fh(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):
    N = integrate.quad(lambda E: DOSfe(E, T, EF), E0, E1, limit = nmaxdiv, epsrel = eps)
    if Debug:
        print("  Ne integ error: {}".format(N[1]))
    return N[0]


def Nh(T, EF, E0, E1):
    N = integrate.quad(lambda E: DOSfh(E, T, EF), E0, E1, limit = nmaxdiv, epsrel = eps)
    if Debug:
        print("  Nh integ error: {}".format(N[1]))
    return N[0]

# ionized donor density
def NDp(EF, T):
    global ND, ED, kB
    return ND * (1.0 - fe(ED, EF, T))

# ionized acceptor density
def NAm(EF, T):
    global NA, EA, kB
    return NA * fe(EA, EF, T)

# 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


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("  Integration E range: {} - {} eV, or EF+-{}*kBT eV".format(Einteg0, Einteg1, nrange))
    print("")
    EV, EC = FindBandEdges(E_raw, dos_raw, EF0, Egth)
    Eg = EC - EV
    print("Band edge: EV={}  EC={}  Eg={} eV".format(EV, EC, Eg))
    print("EF at 0 K = ", EF0, " eV")
    Ne0K = Ne(0.0, EF0, Einteg0, EF0)
    print("Ne at 0 K = ", Ne0K, " /u.c.")
    print("")

    xT = []
    yEF = []
    yEint = []
    yEentropy = []
    yEFa = []
    EFprev = EF0
    print ("%5s\t%8s\t%14s" % ("T(K)", "EF(eV)", "Ne(0 K) (check)"))
    for iT in range(nT):
        T = Tmin + iT * Tstep
        kBTe = kB * T / e
        dE = nrange * kBTe

        E0 = EV - dE
        E1 = EC + dE

        if E0 < Emin or Emax < E1:
            print("*** Integration range error: [{}, {}] exceeds DOS E range [{}, {}]".format(E0, E1, Emin, Emax))
            exit()

        Ne0K = Ne(0.0, EF0, E0, EF0)

#        EF = CalEF(T, EFprev, E0, E1, Ne0K)

#初期範囲
        EFmin = EF0 - dEbisec
        EFmax = EF0 + dEbisec
# まず、EFmin,EFmaxにおけるΔQを計算し、それぞれが正・負あるいは負・生となっていることを確認する    
        print("E range:", EF0, E1)
        Nemin = Ne(T, EFmin, EF0, E1)
        Nhmin = Nh(T, EFmin, E0, EF0)
        dQmin = Nemin - Nhmin
        Nemax = Ne(T, EFmax, EF0, E1)
        Nhmax = Nh(T, EFmax, E0, EF0)
        dQmax = Nemax - Nhmax
        print("iT {:03d}: min (EF,Ne,Nh)=({:6.4f},{:10.4g},{:10.4g},{:10.4g})".format(iT, EFmin, Nemin, Nhmin, dQmin))
        print("        max (EF,Ne,Nh)=({:6.4f},{:10.4g},{:10.4g},{:10.4g})".format(EFmax, Nemax, Nhmax, dQmax))
#        dQmin = Ne(EFmin, T) + NAm(EFmin, T) - Nh(EFmin, T) - NDp(EFmin, T)
#        dQmax = Ne(EFmax, T) + NAm(EFmax, T) - Nh(EFmax, T) - NDp(EFmax, T)
#        print("  EFmin = {:12.8f}  dQmin = {:12.4g}".format(EFmin, dQmin))
#        print("  EFmax = {:12.8f}  dQmax = {:12.4g}".format(EFmax, dQmax))
        print("dQ=", dQmin, dQmax)
        print("dQmin*dQmax=", dQmin * dQmax)
        if dQmin * dQmax > 0.0:
            print("Error: Initial Emin and Emax should be chosen as dQmin * dQmax < 0")
            return 0

# 2分法開始
        for i in range(nmaxiter):
            EFhalf = (EFmin + EFmax) / 2.0
            Neh  = Ne(T, EFhalf, EF0, E1)
#            NAmh = NAm(EFhalf, T)
            Nhh  = Nh(T, EFhalf, E0, EF0)
#            NDph = NDp(EFhalf, T)
            dQhalf = Neh - Nhh
            print("  iter {:03d}: half (EF,Ne,Nh)=({:6.4f},{:10.4g},{:10.4g},{:10.4g})".format(i, EFhalf, Neh, Nhh, dQhalf))
#            dQhalf = Neh + NAmh - Nhh - NDph
#            print("  Iter {}: EFhalf = {:12.8f}  dQhalf = {:12.4g}".format(i, EFhalf, dQhalf))
#            print("     Ne={:10.4e}  Nh={:10.4e}  NA-={:10.4e}  ND+={:10.4e}  dQ={:10.4e}".format(Neh, Nhh, NAmh, NDph, dQhalf))
# EFの精度がepsより小さくなったら計算終了
            if abs(EFmin - EFhalf) < epsbisec and abs(EFmax - EFhalf) < epsbisec:
#                    print("  Success: Convergence reached at EF = {}".format(EFhalf))
                    break
            if dQmin * dQhalf < 0.0:
                EFmax = EFhalf
                dQmax = dQhalf
            else:
                EFmin = EFhalf
                dQmin = dQhalf
        else:
            print("  Failed: Convergence did not reach")
            return 0

        EF = EFhalf

        xT.append(T)
        yEF.append(EF)

        EFprev = EF
        print ("%5.0f\t%12.6g" % (T, EF))


#=============================
# 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, yEF, label = 'EF', color = 'black')
    ax1.set_xlabel("T (K)")
    ax1.set_ylabel("EF (eV)")
    ax1.legend()
    ax2.plot(E_raw, dos_raw, label = 'DOS', color = 'black')
    ax2.plot([EF0, EF0], [min(dos_raw), max(dos_raw)], color = 'red', linestyle = 'dashed')
    ax2.set_xlabel("E (eV)")
    ax2.set_ylabel("DOS (states/u.c.)")
    ax2.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()