"""
Bose-Einstein凝縮の計算と分析を行うモジュール

概要:
    このモジュールは、理想ボーズ気体におけるBose-Einstein凝縮現象をシミュレーションし、
    関連する物理量（化学ポテンシャル、粒子数、凝縮粒子数など）を計算およびプロットします。

詳細説明:
    コマンドライン引数により、動作モード（Fs-α関数のプロットまたは化学ポテンシャルと粒子数の計算）と
    計算範囲を設定できます。主要な機能として、Bose-Einstein積分の計算、
    与えられた粒子数密度に対する化学ポテンシャルの決定（二分法を使用）、
    そしてそれらの結果のグラフ表示が含まれています。

関連リンク: :doc:`bose_condensation_usage`
"""

import sys
import time
from math import exp, sqrt
import numpy as np
from scipy import integrate         # 数値積分関数 integrateを読み込む
from scipy import optimize          # newton関数はscipy.optimizeモジュールに入っている
from matplotlib import pyplot as plt

#===================
# 物理定数
#===================
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-1 (J/K)
ME   = 9.1093897e-31    # kg
MP   = 1.6726231e-27    # kg
MN   = 1.67495e-27      # kg

#===================
# デフォルトパラメータ
#===================
m_def = 6.64e-27      # kg, 4He  
N_def = 2.18e22       # cm^-3, Liq He

# Temperature range
Tmin_def  =  3.00     # K
Tmax_def  =  4.00
Tstep_def =  0.01

# alpha = -mu / kB / T range
alphamin_def  = 0.0
alphamax_def  = 2.0
alphastep_def = 0.02

# 収束条件
EPS_DEF   = 1.0e-12
NMAXITER = 100

def Gamma(sigma: float) -> float:
    """
    ガンマ関数を計算します。

    概要: 階乗関数の実数・複素数への一般化であるガンマ関数を再帰的に計算します。
    詳細説明:
        この実装は再帰を使用し、`sigma`が1や0.5に近い特殊なケースを処理します。
        `sigma`が0.5未満の場合、関数はエラーメッセージを表示してプログラムを終了します。
    :param sigma: float: ガンマ関数を計算する入力値。
    :returns: float: 与えられた`sigma`に対するガンマ関数の計算結果。
    """
    if abs(sigma - 1.0) < 1.0e-6:
        return 1.0
    if abs(sigma - 0.5) < 1.0e-6:
        return sqrt(PI)
    if sigma < 0.5 - 1.0e-6:
        print("Gamma: Abnormal argument sigma = ", sigma)
        sys.exit()
    return (sigma - 1.0) * Gamma(sigma - 1.0)

def IntegFunc(y: float, sigma: float, alpha: float) -> float:
    """
    Bose-Einstein積分 F_sigma(alpha) の被積分関数を定義します。

    概要: Bose-Einstein分布関数に関連する積分の一部を形成する被積分関数を計算します。
    :param y: float: 積分変数。
    :param sigma: float: 関数の次数。
    :param alpha: float: 無次元パラメータ（-mu/(kBT)）。
    :returns: float: 被積分関数の計算結果。
    """
    if y + alpha > 100.0:
        return 0.0
    # y=0付近での発散を避けるため微小値を加算する等の処理はquadに任せる
    return pow(y, sigma - 1.0) / (exp(y + alpha) - 1.0)

def Fsalpha(sigma: float, alpha: float, Emax: float = 10.0, eps: float = 1.0e-8) -> float:
    """
    Bose-Einstein積分 F_sigma(alpha) を数値的に計算します。

    概要: Bose-Einstein関数 F_sigma(alpha) を数値積分を用いて計算します。
    """
    # 0付近の特異性に対応するため points=[0] を指定する場合もあるが、quadのデフォルトで概ね動作する
    ret = integrate.quad(lambda y: IntegFunc(y, sigma, alpha), 0.0, Emax, epsrel = eps)
    return 1.0 / Gamma(sigma) * ret[0]

def Ne_func(EF: float, T: float, mass: float, Emax: float = 10.0, eps: float = 1.0e-8) -> float:
    """
    特定の化学ポテンシャル (EF) と温度 (T) におけるボーズ粒子の密度を計算します。
    """
    lambdaT = H / sqrt(2.0 * PI * mass * KB * T)
    alpha = -EF / (KB * T / E)
    # 粒子数密度を [cm^-3] 単位に変換 (* 1.0e-6)
    density = 1.0 / pow(lambdaT, 3.0) * Fsalpha(1.5, alpha, Emax = Emax, eps = eps) * 1.0e-6
    return density

def CalEF(T: float, target_N: float, mass: float, EFmin: float, EFmax: float, eps_val: float) -> float:
    """
    二分法を用いて、指定された温度と粒子数密度に対する化学ポテンシャルを決定します。
    """
    dQmin = Ne_func(EFmin, T, mass) - target_N
    dQmax = Ne_func(EFmax, T, mass) - target_N

    if dQmin * dQmax > 0.0:
        print(f"Error: Initial range [{EFmin}, {EFmax}] does not bracket the root (dQmin*dQmax > 0)")
        return 0

    EFhalf = (EFmin + EFmax) / 2.0
    for i in range(NMAXITER):
        EFhalf = (EFmin + EFmax) / 2.0
        dQhalf = Ne_func(EFhalf, T, mass) - target_N
        if abs(EFmin - EFhalf) < eps_val and abs(EFmax - EFhalf) < eps_val:
            break
        if dQmin * dQhalf < 0.0:
            EFmax = EFhalf
            dQmax = dQhalf
        else:
            EFmin = EFhalf
            dQmin = dQhalf
    else:
        print("Warning: Convergence did not reach max iterations")
    return EFhalf

def ExecFs(params):
    """Fs-α曲線と粒子数密度Nのプロットを実行します。"""
    Tmin, Tmax, Tstep = params['Tmin'], params['Tmax'], params['Tstep']
    alphamin, alphamax, alphastep = params['alphamin'], params['alphamax'], params['alphastep']
    mass, target_N, zeta32 = params['mass'], params['target_N'], params['zeta32']

    nT = int((Tmax - Tmin) / Tstep + 1.00001)
    nalpha = int((alphamax - alphamin) / alphastep + 1.00001)

    fig, ax1 = plt.subplots()
    ax1.set_xlabel("alpha = -mu / kB / T")
    ax1.set_ylabel("Ne (cm-3)")

    for i in range(nT):
        T = Tmax - i * Tstep
        lambdaT = H / sqrt(2.0 * PI * mass * KB * T)
        Nmax = 1.0 / pow(lambdaT, 3.0) * zeta32 * 1.0e-6
        print(f"T = {T:g} K  lambda_T = {lambdaT * 1e9:g} nm  Nmax = {Nmax:g} cm-3")

        xalpha, yNe = [], []
        for ia in range(nalpha):
            alpha = alphamin + ia * alphastep
            EF = -alpha * KB * T / E
            ne = Ne_func(EF, T, mass) 
            xalpha.append(alpha)
            yNe.append(ne)
        
        ax1.plot(xalpha, yNe, label=f'Ne({T:g} K)', linewidth=0.8)

    ax1.plot([alphamin, alphamax], [target_N, target_N], color='red', linewidth=1.0, linestyle='dashed', label='Target N')
    ax1.set_xlim([alphamin, alphamax])
    ax1.legend()
    plt.title("Particle Density vs Alpha")
    plt.show(block=False)
    input("Press ENTER to exit>>")

def ExecMu(params):
    """化学ポテンシャル (mu) および関連するパラメータの計算とプロットを実行します。"""
    Tmin, Tmax, Tstep = params['Tmin'], params['Tmax'], params['Tstep']
    mass, target_N, zeta32, eps_val = params['mass'], params['target_N'], params['zeta32'], params['eps']

    Tc = H * H / 2.0 / PI / mass / KB * pow(target_N * 1.0e6 / zeta32, 2.0/3.0)
    nT = int((Tmax - Tmin) / Tstep + 1.00001)

    xT, yEF, yEFapprox = [], [], []
    print("\n   T(K)   \tlambda_T(nm)\t   EF(eV)\t   Fapprox(eV)\t   N'(cm-3)\t   Nmax(cm-3)\t   N0(cm-3)")
    
    for i in range(nT):
        T = Tmax - i * Tstep
        lambdaT = H / sqrt(2.0 * PI * mass * KB * T)
        Nmax = 1.0 / pow(lambdaT, 3.0) * zeta32 * 1.0e-6
        
        if T < Tc:
            EF = 0.0
            EFapprox = 0.0
            Ncal = target_N * pow(T / Tc, 1.5)
            N0 = target_N - Ncal
            print(f"{T:8.5f}\t{lambdaT*1e9:8.4g}\t{EF:12.4g}\t{EFapprox:12.4g}\t{Ncal:12.4e}\t{Nmax:12.4e}\t{N0:12.4e}")
        else:
            EFmin, EFmax = -2.0, -1.0e-100
            EF = CalEF(T, target_N, mass, EFmin, EFmax, eps_val)
            t_rel = (T - Tc) / Tc
            Aapprox = 9.0 / 16.0 / PI * zeta32 * zeta32 * t_rel * t_rel
            EFapprox = -Aapprox * KB * T / E
            Ncal = Ne_func(EF, T, mass)
            print(f"{T:8.5f}\t{lambdaT*1e9:8.4g}\t{EF:12.4g}\t{EFapprox:12.4g}\t{Ncal:12.4e}\t{Nmax:12.4e}")
            
        xT.append(T)
        yEF.append(-EF / (KB * T / E))
        yEFapprox.append(-EFapprox / (KB * T / E))

    fig, ax1 = plt.subplots()
    ax1.plot(xT, yEF, label='alpha (calculated)')
    ax1.plot(xT, yEFapprox, label='alpha (approx)', linestyle='--', color='red')
    ax1.set_xlabel("T (K)")
    ax1.set_ylabel("alpha = -mu/(kBT)")
    ax1.set_xlim([Tmin, Tmax])
    ax1.legend()
    plt.title("Chemical Potential vs Temperature")
    plt.show(block=False)
    input("Press ENTER to exit>>")

def main():
    """スクリプトの主要な実行フローを制御します。"""
    # 初期値の設定
    params = {
        'mass': m_def, 'target_N': N_def,
        'Tmin': Tmin_def, 'Tmax': Tmax_def, 'Tstep': Tstep_def,
        'alphamin': alphamin_def, 'alphamax': alphamax_def, 'alphastep': alphastep_def,
        'eps': EPS_DEF
    }
    mode = 'Fs'

    # コマンドライン引数の処理
    argv = sys.argv
    if len(argv) > 1:
        mode = argv[1]
        if mode == 'mu':
            if len(argv) >= 3: params['Tmin'] = float(argv[2])
            if len(argv) >= 4: params['Tmax'] = float(argv[3])
            if len(argv) >= 5: params['Tstep'] = float(argv[4])
        elif mode == 'Fs':
            if len(argv) >= 3: params['Tmin'] = float(argv[2])
            if len(argv) >= 4: params['Tmax'] = float(argv[3])
            if len(argv) >= 5: params['Tstep'] = float(argv[4])
            if len(argv) >= 6: params['alphamin'] = float(argv[5])
            if len(argv) >= 7: params['alphamax'] = float(argv[6])
            if len(argv) >= 8: params['alphastep'] = float(argv[7])
    else:
        print("Usage: python bose_condensation.py mu Tmin Tmax Tstep")
        print("Usage: python bose_condensation.py Fs Tmin Tmax Tstep alphamin alphamax alphastep")
        return

    # 基本定数の計算
    params['zeta32'] = Fsalpha(1.5, 0.0)
    print(f"mass = {params['mass']} kg, N = {params['target_N']} cm^-3")
    print(f"G(1.5) = {Gamma(1.5):.6f}, F3/2(0) = {params['zeta32']:.6f}")
    
    Tc = H * H / 2.0 / PI / params['mass'] / KB * pow(params['target_N']*1.0e6 / params['zeta32'], 2.0/3.0)
    print(f"Calculated Tc = {Tc:.6f} K\n")

    if mode == 'Fs':
        ExecFs(params)
    elif mode == 'mu':
        ExecMu(params)

if __name__ == '__main__':
    main()