import numpy as np
import matplotlib.pyplot as plt
from scipy.constants import hbar, m_e, pi, k, e

# -----------------------------
# パラメータ設定
# -----------------------------
meff = 0.2               # 有効質量 (単位 m_e)
b = -1.0e-58             # quartic 項の係数 [J·m^4]
T = 300                  # 温度 [K]

Emin = 0.0
Emax = 1.0 * 1.6e-19   # 1 eV
nE = 400
Egrid = np.linspace(Emin, Emax, nE)
dE = Egrid[1]-Egrid[0]

# -----------------------------
# DOS(E)
# -----------------------------
def DOS_analytic(E, meff, b):
    a_coeff = hbar**2 / (2 * meff * m_e)
    if b == 0.0:  # parabolic limit
        return (1/(2*pi**2)) * (2*meff*m_e/hbar**2)**1.5 * np.sqrt(np.maximum(E,0))
    disc = a_coeff**2 + 4*b*E
    k2 = (-a_coeff + np.sqrt(np.maximum(disc,0))) / (2*b)
    k2 = np.where(k2 > 0, k2, 0.0)
    k = np.sqrt(k2)
    denom = a_coeff + 2*b*k2
    return (1.0/(4*pi**2)) * np.divide(k, denom, out=np.zeros_like(k), where=denom>0)

# 群速度
def v_E(E, meff, b):
    a_coeff = hbar**2 / (2 * meff * m_e)
    if b == 0.0:
        k = np.sqrt(2*meff*m_e*E)/hbar
        return hbar*k/(meff*m_e)
    disc = a_coeff**2 + 4*b*E
    k2 = (-a_coeff + np.sqrt(np.maximum(disc,0))) / (2*b)
    k2 = np.where(k2 > 0, k2, 0.0)
    k = np.sqrt(k2)
    dEdk = 2*a_coeff*k + 4*b*k**3
    return dEdk/hbar

# 輸送DOS
def sigma_E(E, meff, b):
    return DOS_analytic(E, meff, b) * v_E(E, meff, b)**2

# Seebeck係数（Mott式）
def Seebeck(EF, meff, b, T):
    sigE = sigma_E(Egrid, meff, b)
    logsig = np.log(np.where(sigE > 0, sigE, 1e-300))
    dlogsigma_dE = np.gradient(logsig, dE)
    idx = np.argmin(np.abs(Egrid - EF))
    return -(pi**2 * k**2 * T)/(3*e) * dlogsigma_dE[idx]

# キャリア濃度
def carrier_density(EF, meff, b):
    mask = Egrid <= EF
    return np.trapz(DOS_analytic(Egrid[mask], meff, b), Egrid[mask])

# -----------------------------
# EF を走査して S と n を計算
# -----------------------------
EF_list = np.linspace(0.01*1.6e-19, 1.0*1.6e-19, 100)  # J
S_nonpar, S_para = [], []
n_nonpar, n_para = [], []

for EF in EF_list:
    S_nonpar.append(Seebeck(EF, meff, b, T))
    S_para.append(Seebeck(EF, meff, 0.0, T))  # parabolic
    n_nonpar.append(carrier_density(EF, meff, b))
    n_para.append(carrier_density(EF, meff, 0.0))

S_nonpar, S_para = np.array(S_nonpar), np.array(S_para)
n_nonpar, n_para = np.array(n_nonpar), np.array(n_para)

# -----------------------------
# プロット (S vs n)
# -----------------------------
plt.figure(figsize=(6,4))
plt.plot(n_para/1e27, S_para*1e6, label="Parabolic band (b=0)")
plt.plot(n_nonpar/1e27, S_nonpar*1e6, label=f"Non-parabolic (b={b:.1e})")
plt.xlabel("Carrier density n [10^27 m^-3]")
plt.ylabel("Seebeck coefficient S [µV/K]")
plt.title(f"S vs Carrier density (meff={meff} m_e, T={T} K)")
plt.legend()
plt.tight_layout()
plt.show()