"""
概要:
    バンド構造から有効質量と運動量空間の波数kの関係を計算します。
詳細説明:
    このスクリプトは、CSVファイルからバンド構造データ（運動量空間の波数kとエネルギーE）を読み込みます。
    読み込んだデータを用いて、エネルギーの運動量空間の波数kに対する二階微分を計算し、そこから有効質量を導出します。
    有効質量の計算では、異なる差分幅（nskip）を使用して二階微分を評価し、結果をプロットに表示します。
    計算された有効質量の逆数が特定の閾値を超えた場合、プロット上でデータの連続性を分断する処理が含まれます。
    最終的なプロットは、有効質量と運動量空間の波数kの関係を示します。
主な機能:
    - CSVファイルからのバンド構造データの読み込み
    - エネルギーの二階微分に基づいた有効質量の計算
    - 異なる差分幅での有効質量計算の比較
    - matplotlibによる結果のグラフ表示
関連リンク:
    EffectiveMass_usage
"""
import sys
import numpy as np
from numpy import sqrt, exp, sin, cos, tan, pi
import numpy.linalg as LA 
import csv
from matplotlib import pyplot as plt


#===================================
# physical constants
#===================================
pi   = 3.14159265358979323846
pi2  = 2.0 * pi
h    = 6.6260755e-34    # Js";
hbar = 1.05459e-34      # "Js";
c    = 2.99792458e8     # m/s";
e    = 1.60218e-19      # C";
e0   = 8.854418782e-12; # C<sup>2</sup>N<sup>-1</sup>m<sup>-2</sup>";
kB   = 1.380658e-23     # JK<sup>-1</sup>";
me   = 9.1093897e-31    # kg";
R    = 8.314462618      # J/K/mol
a0   = 5.29177e-11      # m";


#===================================
# parameters
#===================================
a = 4.0e-10 #m

infile = 'band.csv'
cutline = 1

#===================================
# figure configuration
#===================================
fontsize        = 12
legend_fontsize = 8


#=============================
# other functions
#=============================
def savecsv(outfile, header, datalist):
    """
    概要:
        リストのリスト形式のデータをCSVファイルに保存します。
    詳細説明:
        指定されたファイル名でCSVファイルを開き、ヘッダー行を書き込んだ後、datalistの各要素を列として、
        行ごとにデータを書き込みます。ファイル書き込みに失敗した場合はエラーメッセージを表示します。
    引数:
        :param outfile: 出力するCSVファイルのパスとファイル名。
        :type outfile: str
        :param header: CSVファイルのヘッダー行の文字列のリスト。
        :type header: list of str
        :param datalist: 出力するデータを含むリストのリスト。各内部リストは1つの列に対応します。
                         例: [[col1_data1, col1_data2], [col2_data1, col2_data2]]
    戻り値:
        :returns: なし
        :rtype: None
    """
    try: 
        print("Write to [{}]".format(outfile))
        f = open(outfile, 'w')
    except:
#    except IOError:
        print("Error: Can not write to [{}]".format(outfile))
    else:
        fout = csv.writer(f, lineterminator='\n')
        fout.writerow(header)
#        fout.writerows(data)
        for i in range(0, len(datalist[0])):
            a = []
            for j in range(len(datalist)):
                a.append(datalist[j][i])
            fout.writerow(a)
        f.close()


def read_csv(fname):
    """
    概要:
        指定されたCSVファイルから数値データを読み込みます。
    詳細説明:
        CSVファイルをオープンし、最初の行をラベルとして読み込みます。
        残りの行から2列のデータを浮動小数点数として読み取り、それぞれのリストに格納します。
        数値変換に失敗した行は警告メッセージを表示してスキップします。
    引数:
        :param fname: 読み込むCSVファイルのパスとファイル名。
        :type fname: str
    戻り値:
        :returns: 1列目のラベル文字列、2列目のラベル文字列、1列目のデータリスト、2列目のデータリストのタプル。
        :rtype: tuple of (str, str, list of float, list of float)
    """
    x = []
    y = []
    with open(fname) as f:
        fin = csv.reader(f)
        labels = next(fin)
        xlabel = labels[0]
        ylabel = labels[1]
        for row in fin:
            try:
                x.append(float(row[0]))
                y.append(float(row[1]))
            except:
                print("Warning: Invalid float data [{}] or [{}]".format(row[0], row[1]))

    return xlabel, ylabel, x, y


def main():
    """
    概要:
        有効質量と運動量空間の波数kの関係を計算し、結果をプロットします。
    詳細説明:
        `infile` で指定されたCSVファイルからバンド構造データ（運動量空間の波数kとエネルギーE）を読み込みます。
        波数kの差分dkとエネルギーEを用いて、エネルギーの運動量空間の波数kに対する二階微分を計算し、
        そこから有効質量を導出します。
        異なる差分幅（nskip = 1とnskip = 4）で二階微分を計算し、それぞれ有効質量を求めます。
        有効質量の逆数 abs(minv) が 1.0e20 以下の場合は、有効質量のデータとして None を格納し、
        プロット上でデータの連続性を分断します (cutline が 1 の場合)。
        また、有効質量の符号が反転した場合もプロットを分断します。
        最終的に、計算された有効質量を電子の静止質量 me で正規化し、
        matplotlib を用いて運動量空間の波数kに対する正規化有効質量をプロットし、表示します。
    引数:
        :param なし: この関数は引数を取りません。
    戻り値:
        :returns: なし
        :rtype: None
    """
    klabel, Elabel, k, E = read_csv(infile)
    print("k=", k)
    print("E=", E)

    nk = len(k)
    dk = k[1] - k[0]

    km = hbar * hbar * (pi2 / a)**2.0
    nskip = 1
    xk  = []
    ymc = []
    signprev = None
    print("")
    print("Differentiate for nskip = ", nskip)
    print("i    E[i-1]  E[i]    E[i+1]  1.0/m*")
    for i in range(nskip, nk - nskip, nskip):
        d2Edk2c = (E[i+nskip] + E[i-nskip] - 2 * E[i]) * e / pow(nskip * dk, 2.0)
        minv = d2Edk2c / km
        print(i, E[i-1], E[i], E[i+1], minv)
        if abs(minv) <= 1.0e20:  # << 1.0/me ~ 1e30 
            if cutline:
                xk.append(k[i])
                ymc.append(None)
            signprev = -signprev
            continue
        else:
            m = km / d2Edk2c

        if signprev is None:
            signprev = m
        elif signprev * m < 0.0:
            if cutline:
                xk.append(k[i])
                ymc.append(None)
            signprev = m

        xk.append(k[i])
        ymc.append(m / me)

    nskip = 4
    xk2  = []
    ymc2 = []
    signprev = None
    print("")
    print("Differentiate for nskip = ", nskip)
    print("i    E[i-1]  E[i]    E[i+1]  1.0/m*")
    for i in range(nskip, nk - nskip, nskip):
        d2Edk2c = (E[i+nskip] + E[i-nskip] - 2 * E[i]) * e / pow(nskip * dk, 2.0)
        minv = d2Edk2c / km
        print(i, E[i-1], E[i], E[i+1], minv)
        if abs(minv) <= 1.0e20:  # << 1.0/me ~ 1e30 
            if cutline:
                xk2.append(k[i])
                ymc2.append(None)
            signprev = -signprev
            continue
        else:
            m = km / d2Edk2c

        if signprev is None:
            signprev = m
        elif signprev * m < 0.0:
            if cutline:
                xk2.append(k[i])
                ymc2.append(None)
            signprev = m

        xk2.append(k[i])
        ymc2.append(m / me)

    plt.plot(xk,  ymc,  linewidth = 0.5, marker = 'o', markersize = 1.0, label = 'nskip = 1')
    plt.plot(xk2, ymc2, linewidth = 0.5, marker = 'o', markersize = 1.0, label = 'nskip = 2')
    plt.xlabel(klabel)
    plt.ylabel("m$_e$ / m$_e^0$")
    plt.xlim([-0.5, 0.5])
#    plt.ylim([-0.5, 0.5])
    plt.tight_layout()
    
    plt.pause(0.1)
    print("Press ENTER to exit>>", end = '')
    input()


if __name__ == "__main__":
    main()