"""
Seebeck_scan.py

概要:
このスクリプトは、フェルミ準位、キャリア濃度、温度、有効質量を様々にスキャンしながら、熱電材料のゼーベック係数を計算しプロットします。

詳細説明:
散乱因子 `r` の異なる値に対して、以下のスキャンモードをサポートしています。
-   EF (フェルミ準位): フェルミ準位をスキャンし、ゼーベック係数を計算します。
-   Ne (キャリア濃度): キャリア濃度をスキャンし、ゼーベック係数を計算します。
-   T (温度): 温度をスキャンし、ゼーベック係数を計算します。このモードでは、フェルミ準位、キャリア濃度、または還元フェルミ準位のいずれかを固定できます。
-   m_eff (有効質量): 有効質量をスキャンし、ゼーベック係数を計算します。このモードでも、フェルミ準位、キャリア濃度、または還元フェルミ準位のいずれかを固定できます。

計算はFermi–Dirac積分に基づき、キャリアタイプ（電子または正孔）を指定できます。
結果はCSVファイルとして保存することも可能です。

散乱因子 r はキャリアの緩和時間 `tau(E)` がエネルギー `E` に対して `tau(E)∝E^(r-1/2)` で定義されます。

実行例:

# EF scan
python Seebeck_scan.py --mode EF --T 300 --m_eff 1.0 --scan-min -0.2 --scan-max 0.5 --n 121

# Ne scan
python Seebeck_scan.py --mode Ne --T 300 --m_eff 1.0 --scan-min 1e15 --scan-max 1e21 --n 121 --logscan

# T scan, EF固定
python Seebeck_scan.py --mode T --fixed EF --EF 0.10 --m_eff 1.0 --scan-min 100 --scan-max 800 --n 121

# T scan, Ne固定
python Seebeck_scan.py --mode T --fixed Ne --Ne 1e18 --m_eff 1.0 --scan-min 100 --scan-max 800 --n 121

# m_eff scan, Ne固定
python Seebeck_scan.py --mode m_eff --fixed Ne --Ne 1e18 --T 300 --scan-min 0.1 --scan-max 3.0 --n 121

# CSV保存
python Seebeck_scan.py --mode Ne --csv result.csv

関連リンク:
:doc:`Seebeck_scan_usage`
"""
import argparse
import csv
from functools import lru_cache

import matplotlib.pyplot as plt
import numpy as np
from scipy.integrate import quad
from scipy.optimize import brentq
from scipy.special import expit, gamma

try:
    from tklib.tksci.tktransport.tkfermi_integral import FermiIntegral_fast
except:
    FermiIntegral_fast = None

# =========================
# constants (SI)
# =========================
kB_SI = 1.380649e-23           # J/K
e_SI  = 1.602176634e-19        # C
h_SI  = 6.62607015e-34         # J s
m0_SI = 9.1093837015e-31       # kg
kB_eV = 8.617333262145e-5      # eV/K


# ============================================================
# Fermi–Dirac integral (NO 1/Gamma(j+1) prefactor)
#   Fj(j,xi) = ∫_0^∞ x^j / (1+exp(x-xi)) dx
# ============================================================
@lru_cache(maxsize=20000)
def Fj(j, xi):
    """
    フェルミ・ディラック積分を計算します（ガンマ関数による前因子なし）。

    概要:
    Fj(j, xi) = ∫_0^∞ x^j / (1+exp(x-xi)) dx を数値的に評価します。

    詳細説明:
    この関数は、引数 j と xi に対してフェルミ・ディラック積分を計算します。
    標準的な定義に含まれる 1/Gamma(j+1) の前因子は含まれません。
    計算結果は `@lru_cache` デコレータによってキャッシュされ、同じ引数での再計算を防ぎ、
    パフォーマンスを向上させます。
    j が -1.0 以下の場合、ValueError を発生させます。

    :param j: float; フェルミ・ディラック積分の次数。j > -1 である必要があります。
    :param xi: float; 還元フェルミ準位 (xi = E_F / (k_B T))。
    :returns: float; 計算されたフェルミ・ディラック積分の値。
    """
    j = float(j)
    xi = float(np.round(xi, 12))
    if j <= -1.0:
        raise ValueError(f"Fj(j, xi) requires j > -1, got j={j}")

    integrand = lambda x: (x**j) * expit(xi - x)
    val, _ = quad(integrand, 0.0, np.inf, epsabs=1e-10, epsrel=1e-10, limit=400)
    return val

fermi_integral_type = "build-in"
fermi_integral_func = Fj
if FermiIntegral_fast:
    fermi_integral_type = "tkfermi_integral.FermiIntegral_fast"
    fermi_integral_func = lambda j, xj: FermiIntegral_fast(xj, j, epsabs=1.0e-10, epsrel=1.0e-8, limit=50, use_cache = True)

def Nc_3D(T, m_eff=1.0):
    """
    3次元における有効状態密度 Nc を計算します。

    概要:
    温度 T と有効質量 m_eff に基づいて、3次元バンドの有効状態密度 Nc を計算します。

    :param T: float; 温度 [K]。
    :param m_eff: float; 有効質量比 m*/m0。m0 は自由電子質量です。
    :returns: float; 計算された有効状態密度 Nc [m^-3]。
    """
    mstar = m_eff * m0_SI
    return 2.0 * (2.0 * np.pi * mstar * kB_SI * T / (h_SI**2)) ** 1.5


def electron_density_from_xi(xi, T, m_eff=1.0):
    """
    還元フェルミ準位 xi から電子濃度を計算します。

    概要:
    還元フェルミ準位 xi、温度 T、有効質量 m_eff を用いて電子濃度 n を計算します。

    詳細説明:
    電子濃度 n は以下の式で与えられます:
    n = Nc * F_{1/2}^{std}(xi) = Nc * Fj(1/2, xi) / Gamma(3/2)
    ここで、Nc は有効状態密度、Fj(1/2, xi) はガンマ関数による前因子を含まないフェルミ・ディラック積分です。

    :param xi: float; 還元フェルミ準位。
    :param T: float; 温度 [K]。
    :param m_eff: float; 有効質量比 m*/m0。
    :returns: float; 計算された電子濃度 n [m^-3]。
    """
    Nc = Nc_3D(T, m_eff=m_eff)
    F12_std = Fj(0.5, xi) / gamma(1.5)
    return Nc * F12_std


def xi_from_electron_density(n_m3, T, m_eff=1.0, xi_low=-40.0, xi_high=120.0):
    """
    与えられた電子濃度から還元フェルミ準位 xi を計算します。

    概要:
    電子濃度 n_m3、温度 T、有効質量 m_eff に対応する還元フェルミ準位 xi を見つけます。

    詳細説明:
    この関数は、`electron_density_from_xi` 関数を解いて xi を見つけるために、
    Brentのアルゴリズム (`scipy.optimize.brentq`) を使用します。
    初期の探索範囲 `[xi_low, xi_high]` で解が見つからない場合、
    自動的に探索範囲を拡張しようと試みます。
    電子濃度が0以下の場合、ValueError を発生させます。

    :param n_m3: float; 電子濃度 [m^-3]。正の値である必要があります。
    :param T: float; 温度 [K]。
    :param m_eff: float; 有効質量比 m*/m0。
    :param xi_low: float; brentqの初期探索範囲の下限。
    :param xi_high: float; brentqの初期探索範囲の上限。
    :returns: float; 計算された還元フェルミ準位 xi。
    :raises ValueError: 電子濃度が正でない場合。
    :raises RuntimeError: 指定された電子濃度に対するxiの探索範囲を見つけられなかった場合。
    """
    if n_m3 <= 0.0:
        raise ValueError(f"Electron density must be positive, got {n_m3}")

    def f(xi):
        return electron_density_from_xi(xi, T, m_eff) - n_m3

    fl = f(xi_low)
    fh = f(xi_high)

    # 探索範囲を動的に拡張
    n_expand = 0
    while fl > 0.0 and n_expand < 20:
        xi_low -= 20.0
        fl = f(xi_low)
        n_expand += 1

    n_expand = 0
    while fh < 0.0 and n_expand < 20:
        xi_high += 20.0
        fh = f(xi_high)
        n_expand += 1

    if fl * fh > 0.0:
        raise RuntimeError(
            "Failed to bracket xi for the requested electron density. "
            f"n={n_m3:.6e} m^-3, T={T}, m_eff={m_eff}, "
            f"f({xi_low})={fl:.6e}, f({xi_high})={fh:.6e}"
        )

    return brentq(f, xi_low, xi_high, xtol=1e-12, rtol=1e-12, maxiter=200)


def seebeck_from_xi(xi, r, carrier="electron"):
    """
    還元フェルミ準位 xi からゼーベック係数を計算します。

    概要:
    還元フェルミ準位 xi、散乱因子 r、キャリアタイプに基づいてゼーベック係数 S を計算します。

    詳細説明:
    ゼーベック係数は以下の式で計算されます。
    S = (k_B / q) * [((r + 2) / (r + 1)) * (F_{r+1}(xi) / F_r(xi)) - xi]
    ここで q はキャリアの電荷（電子の場合は -e_SI、正孔の場合は +e_SI）です。
    散乱因子 r は -1.0 より大きい必要があります。

    :param xi: float; 還元フェルミ準位。
    :param r: float; 散乱因子。r > -1 である必要があります。
    :param carrier: str; キャリアのタイプ ("electron" または "hole")。
    :returns: float; 計算されたゼーベック係数 S [V/K]。
    :raises ValueError: キャリアタイプが 'electron' または 'hole' ではない場合、または r が -1.0 以下の場合。
    """
    if carrier.lower().startswith("e"):
        q = -e_SI
    elif carrier.lower().startswith("h"):
        q = +e_SI
    else:
        raise ValueError("carrier must be 'electron' or 'hole'")

    if r <= -1.0:
        raise ValueError(f"r must be > -1, got r={r}")

    Fr = Fj(r, xi)
    Fr1 = Fj(r + 1.0, xi)
    bracket = ((r + 2.0) / (r + 1.0)) * (Fr1 / Fr) - xi
    return (kB_SI / q) * bracket


def parse_r_list(text):
    """
    カンマ区切りの文字列から散乱因子 r のリストをパースします。

    概要:
    入力文字列を解析し、浮動小数点数の散乱因子 r のリストを返します。

    :param text: str; カンマで区切られた散乱因子の文字列（例: "2,1.5,1"）。
    :returns: list[float]; パースされた散乱因子 r のリスト。
    :raises ValueError: r のリストが空の場合、またはリスト内のいずれかの r が -1.0 以下の場合。
    """
    vals = [float(x.strip()) for x in text.split(",") if x.strip()]
    if not vals:
        raise ValueError("r-list is empty")
    for r in vals:
        if r <= -1.0:
            raise ValueError(f"Each r must be > -1, got {r}")
    return vals


def default_scan_settings(mode):
    """
    指定されたスキャンモードのデフォルト設定を返します。

    概要:
    スキャンモードに基づいて、スキャン範囲の最小値、最大値、点数、および対数スケールフラグのデフォルト値を決定します。

    :param mode: str; スキャンモード ("EF", "Ne", "T", "m_eff")。
    :returns: tuple[float, float, int, bool]; (scan_min, scan_max, npts, logscan) のタプル。
    :raises ValueError: 未知のスキャンモードが指定された場合。
    """
    if mode == "EF":
        return -0.2, 0.5, 121, False
    if mode == "Ne":
        return 1e15, 1e21, 121, True
    if mode == "T":
        return 100.0, 800.0, 121, False
    if mode == "m_eff":
        return 0.1, 3.0, 121, False
    raise ValueError(f"Unknown mode: {mode}")


def make_scan_values(mode, scan_min, scan_max, npts, logscan):
    """
    スキャンする値の配列を生成します。

    概要:
    指定された範囲 (`scan_min`, `scan_max`) と点数 `npts` に基づいて、
    線形または対数スケールのスキャン値配列を生成します。

    :param mode: str; スキャンモード。ログスキャン時に影響することがあるが、主にエラーメッセージ用。
    :param scan_min: float; スキャン範囲の最小値。
    :param scan_max: float; スキャン範囲の最大値。
    :param npts: int; スキャン点の数。
    :param logscan: bool; スキャン値を対数スケールで生成するかどうか。
    :returns: numpy.ndarray; 生成されたスキャン値の配列。
    :raises ValueError: 対数スキャンで `scan_min` または `scan_max` が0以下の場合。
    """
    if logscan:
        if scan_min <= 0.0 or scan_max <= 0.0:
            raise ValueError("log scan requires scan_min > 0 and scan_max > 0")
        return np.logspace(np.log10(scan_min), np.log10(scan_max), npts)
    return np.linspace(scan_min, scan_max, npts)


def resolve_xi(mode, scan_value, T0, m_eff0, EF0_eV, Ne0_cm3, xi0, fixed):
    """
    現在のスキャン値とモードに基づいて、還元フェルミ準位 (xi)、フェルミ準位 (EF_eV)、
    電子濃度 (Ne_cm3)、温度 (T_K)、および有効質量 (m_eff) を解決します。

    概要:
    様々なスキャンモードと固定条件を考慮して、関連する物理量を計算します。

    詳細説明:
    この関数は、`mode` 引数に応じて、`scan_value` を基に物理量を計算します。
    - `mode` が "EF" または "Ne" の場合、`scan_value` はそれぞれEFまたはNeとして扱われ、Tとm_effは固定値 (`T0`, `m_eff0`) を使用します。
    - `mode` が "T" または "m_eff" の場合、`scan_value` はそれぞれTまたはm_effとして扱われ、`fixed` 引数（"EF", "Ne", "xi"）に基づいて他の量が計算されます。

    :param mode: str; 現在のスキャンモード ("EF", "Ne", "T", "m_eff")。
    :param scan_value: float; 現在のスキャン対象の値。
    :param T0: float; 基準温度 [K]。
    :param m_eff0: float; 基準有効質量比 m*/m0。
    :param EF0_eV: float; 基準フェルミ準位 [eV]。
    :param Ne0_cm3: float; 基準電子濃度 [cm^-3]。
    :param xi0: float; 基準還元フェルミ準位。
    :param fixed: str; `mode` が "T" または "m_eff" の場合に固定する変数 ("EF", "Ne", "xi")。
    :returns: tuple[float, float, float, float, float];
              (xi, EF_eV, Ne_cm3, T_K, m_eff) のタプル。
    :raises ValueError: 未知のスキャンモードまたは固定条件が指定された場合。
    """
    if mode == "EF":
        T = T0
        m_eff = m_eff0
        EF_eV = scan_value
        xi = EF_eV / (kB_eV * T)
    elif mode == "Ne":
        T = T0
        m_eff = m_eff0
        Ne_cm3 = scan_value
        xi = xi_from_electron_density(Ne_cm3 * 1e6, T=T, m_eff=m_eff)
        EF_eV = xi * kB_eV * T
    elif mode == "T":
        T = scan_value
        m_eff = m_eff0
        if fixed == "EF":
            EF_eV = EF0_eV
            xi = EF_eV / (kB_eV * T)
        elif fixed == "Ne":
            xi = xi_from_electron_density(Ne0_cm3 * 1e6, T=T, m_eff=m_eff)
            EF_eV = xi * kB_eV * T
        elif fixed == "xi":
            xi = xi0
            EF_eV = xi * kB_eV * T
        else:
            raise ValueError(f"Unknown fixed condition: {fixed}")
    elif mode == "m_eff":
        T = T0
        m_eff = scan_value
        if fixed == "EF":
            EF_eV = EF0_eV
            xi = EF_eV / (kB_eV * T)
        elif fixed == "Ne":
            xi = xi_from_electron_density(Ne0_cm3 * 1e6, T=T, m_eff=m_eff)
            EF_eV = xi * kB_eV * T
        elif fixed == "xi":
            xi = xi0
            EF_eV = xi * kB_eV * T
        else:
            raise ValueError(f"Unknown fixed condition: {fixed}")
    else:
        raise ValueError(f"Unknown mode: {mode}")

    Ne_cm3 = electron_density_from_xi(xi, T=T, m_eff=m_eff) / 1e6
    return xi, EF_eV, Ne_cm3, T, m_eff


def x_label(mode):
    """
    指定されたスキャンモードに対応するX軸ラベル文字列を返します。

    概要:
    プロットのX軸に表示する適切なラベルを返します。

    :param mode: str; スキャンモード ("EF", "Ne", "T", "m_eff")。
    :returns: str; X軸のラベル文字列。
    """
    return {
        "EF": r"$E_F-E_c$ (eV)",
        "Ne": r"$N_e$ (cm$^{-3}$)",
        "T": r"$T$ (K)",
        "m_eff": r"$m^*/m_0$",
    }[mode]


def make_title(mode, carrier, fixed, T0, m_eff0, EF0_eV, Ne0_cm3, xi0):
    """
    プロットのタイトル文字列を生成します。

    概要:
    スキャンモード、キャリアタイプ、固定条件などのパラメータに基づいて、グラフのタイトルを作成します。

    :param mode: str; スキャンモード ("EF", "Ne", "T", "m_eff")。
    :param carrier: str; キャリアタイプ ("electron" または "hole")。
    :param fixed: str; 固定された変数 ("EF", "Ne", "xi")。
    :param T0: float; 基準温度 [K]。
    :param m_eff0: float; 基準有効質量比 m*/m0。
    :param EF0_eV: float; 基準フェルミ準位 [eV]。
    :param Ne0_cm3: float; 基準電子濃度 [cm^-3]。
    :param xi0: float; 基準還元フェルミ準位。
    :returns: str; 生成されたタイトル文字列。
    """
    base = f"Seebeck coefficient scan ({mode}, {carrier})"
    if mode in ["EF", "Ne"]:
        return base + f"\nT={T0:g} K, m*/m0={m_eff0:g}"
    if fixed == "EF":
        fixed_text = f"fixed EF={EF0_eV:g} eV"
    elif fixed == "Ne":
        fixed_text = f"fixed Ne={Ne0_cm3:.3e} cm^-3"
    else:
        fixed_text = f"fixed xi={xi0:g}"
    return base + f"\n{fixed_text}, T={T0:g} K, m*/m0={m_eff0:g}"


def save_csv(outfile, rows):
    """
    計算結果をCSVファイルに保存します。

    概要:
    計算されたゼーベック係数と関連する物理量を、指定されたファイル名でCSV形式で保存します。

    詳細説明:
    `outfile` が空文字列の場合、ファイル保存はスキップされます。
    CSVヘッダーは固定されており、以下のフィールドが含まれます:
    "mode", "scan_value", "r", "S_V_per_K", "S_uV_per_K",
    "xi", "EF_eV", "Ne_cm3", "T_K", "m_eff_m0"

    :param outfile: str; 出力CSVファイルのパス。空文字列の場合、ファイルは保存されません。
    :param rows: list[dict]; 各行が辞書形式で表現されたデータのリスト。
    """
    if not outfile:
        return
    fieldnames = [
        "mode", "scan_value", "r", "S_V_per_K", "S_uV_per_K",
        "xi", "EF_eV", "Ne_cm3", "T_K", "m_eff_m0"
    ]
    with open(outfile, "w", newline="", encoding="utf-8-sig") as f:
        writer = csv.DictWriter(f, fieldnames=fieldnames)
        writer.writeheader()
        writer.writerows(rows)


def main():
    """
    ゼーベック係数スキャンプログラムのメインエントリポイント。

    概要:
    コマンドライン引数をパースし、指定されたモードとパラメータに基づいてゼーベック係数を計算し、結果をプロットまたはCSVファイルに保存します。

    詳細説明:
    - `argparse` を使用して、スキャンモード、固定条件、キャリアタイプ、基準温度、有効質量、フェルミ準位、電子濃度、還元フェルミ準位、散乱因子のリスト、スキャン範囲、点数、対数スキャンオプション、X軸スケール、CSV出力ファイル名、プロット表示の有無などの引数を定義します。
    - デフォルトのスキャン設定を適用し、スキャン値の配列を生成します。
    - 指定された散乱因子のリスト `r_list` の各値に対して、スキャン値 (`xvals`) ごとにゼーベック係数を計算します。
    - 各計算結果は `rows` リストに追加され、後でCSVとして保存するために使用されます。
    - 計算結果は `matplotlib` を使用してプロットされ、適切なラベル、タイトル、凡例が設定されます。
    - `--csv` オプションが指定された場合、`save_csv` 関数を呼び出して結果を保存します。
    - `--no-show` オプションが指定されない限り、プロットを表示します。
    """
    parser = argparse.ArgumentParser(
        description="Unified Seebeck scan script for EF / Ne / T / m_eff"
    )
    parser.add_argument("--mode", choices=["EF", "Ne", "T", "m_eff"], default="EF",
                        help="scan target")
    parser.add_argument("--fixed", choices=["EF", "Ne", "xi"], default="EF",
                        help="for mode=T or m_eff, what to keep fixed")
    parser.add_argument("--carrier", choices=["electron", "hole"], default="electron")
    parser.add_argument("--T", type=float, default=300.0,
                        help="reference temperature [K]")
    parser.add_argument("--m_eff", type=float, default=1.0,
                        help="reference effective mass m*/m0")
    parser.add_argument("--EF", type=float, default=0.10,
                        help="reference EF-Ec [eV] (or corresponding band-edge reference)")
    parser.add_argument("--Ne", type=float, default=1e18,
                        help="reference electron density [cm^-3]")
    parser.add_argument("--xi", type=float, default=0.0,
                        help="reference reduced Fermi energy xi")
    parser.add_argument("--r-list", default="2,1.5,1,0.5,0,-0.5",
                        help="comma-separated r list")
    parser.add_argument("--scan-min", type=float, default=None)
    parser.add_argument("--scan-max", type=float, default=None)
    parser.add_argument("--n", type=int, default=None,
                        help="number of scan points")
    parser.add_argument("--logscan", action="store_true",
                        help="use log spacing for scan values")
    parser.add_argument("--xscale", choices=["auto", "linear", "log"], default="auto")
    parser.add_argument("--csv", default="",
                        help="optional CSV output filename")
    parser.add_argument("--no-show", action="store_true",
                        help="do not display plot")
    args = parser.parse_args()

    r_list = parse_r_list(args.r_list)

    dmin, dmax, dn, dlog = default_scan_settings(args.mode)
    scan_min = dmin if args.scan_min is None else args.scan_min
    scan_max = dmax if args.scan_max is None else args.scan_max
    npts = dn if args.n is None else args.n
    logscan = args.logscan or dlog

    if npts < 2:
        raise ValueError("n must be >= 2")
    if scan_max <= scan_min:
        raise ValueError("scan-max must be larger than scan-min")
    if args.T <= 0.0:
        raise ValueError("T must be positive")
    if args.m_eff <= 0.0:
        raise ValueError("m_eff must be positive")
    if args.Ne <= 0.0:
        raise ValueError("Ne must be positive")

    xvals = make_scan_values(args.mode, scan_min, scan_max, npts, logscan)

    rows = []
    plt.figure(figsize=(8, 6))

    for r in r_list:
        yvals_uV = []
        for x in xvals:
            xi, EF_eV, Ne_cm3, T_K, m_eff = resolve_xi(
                mode=args.mode,
                scan_value=float(x),
                T0=args.T,
                m_eff0=args.m_eff,
                EF0_eV=args.EF,
                Ne0_cm3=args.Ne,
                xi0=args.xi,
                fixed=args.fixed,
            )
            S = seebeck_from_xi(xi, r=r, carrier=args.carrier)
            yvals_uV.append(S * 1e6)
            rows.append({
                "mode": args.mode,
                "scan_value": float(x),
                "r": float(r),
                "S_V_per_K": float(S),
                "S_uV_per_K": float(S * 1e6),
                "xi": float(xi),
                "EF_eV": float(EF_eV),
                "Ne_cm3": float(Ne_cm3),
                "T_K": float(T_K),
                "m_eff_m0": float(m_eff),
            })

        plt.plot(xvals, yvals_uV, label=f"r = {r:g}")

    save_csv(args.csv, rows)

    plt.axhline(0.0, linewidth=1.0, alpha=0.4)
    plt.xlabel(x_label(args.mode))
    plt.ylabel(r"$S$ ($\mu$V/K)")
    plt.title(make_title(args.mode, args.carrier, args.fixed, args.T, args.m_eff, args.EF, args.Ne, args.xi))

    xscale = args.xscale
    if xscale == "auto":
        xscale = "log" if (args.mode == "Ne" or logscan) else "linear"
    plt.xscale(xscale)

    plt.grid(True, which="both", alpha=0.25)
    plt.legend()
    plt.tight_layout()

    if not args.no_show:
        plt.show()


if __name__ == "__main__":
    main()