"""
水素原子の波動関数プロットスクリプト

概要:
このスクリプトは、水素原子の電子の波動関数、その絶対値、位相、確率密度分布などを様々な形式で可視化します。
コマンドライン引数でプロットモードと関数タイプを指定し、静的プロット、アニメーション、
等値面、3D散布図など多様な表現で波動関数を分析することができます。

詳細説明:
本スクリプトは、以下の主要な機能を提供します。
- `rplot`: 球面調和関数の実部、虚部、絶対値、絶対値の二乗を3D表面としてプロット。
- `ranim`: `rplot`の3D表面プロットをアニメーション化。
- `dot`, `dotc`: 波動関数の確率密度に比例した3D散布図をプロット（位相による色分け）。
- `dotanim`: `dot`プロットをアニメーション化。
- `iso3d`, `iso3danim`: 波動関数の等値面を3Dでプロットし、アニメーション化も可能。
- `Rnl`: 動径波動関数と動径確率密度分布を2Dでプロット。
- `Ylm`, `Phi`: 球面調和関数の角度依存部分（ThetaとPhi）を2Dでプロット。
- `normalize`: 波動関数の正規化を数値的に検証。
- `torus`, `isosurface3d`, `contour3d`, `3dMC_animation`: その他の汎用的な3Dプロットやアニメーション。

利用可能な関数タイプ (`functype`):
- `1s`, `2s`, `2px`, `3dxy` などの軌道指定。
- `Rnl`, `Ylm`, `Theta`, `Phi`, `normalize` などの波動関数コンポーネント。
- `c` (複素数), `r` (実部), `i` (虚部), `a` (絶対値), `a2` (絶対値の2乗) の表示タイプ。

関連リンク:
- 元記事: https://zenn.dev/shittoku_xxx/articles/13afd6fdfac44e
"""

import sys
import argparse
import numpy as np
from numpy import abs, angle, exp, sqrt, cos, sin, arccos, arctan2, pi, min, max
from scipy.integrate import quad
from scipy.constants import physical_constants
from scipy.special import sph_harm, lpmv, assoc_laguerre

import matplotlib
import matplotlib.pyplot as plt
from matplotlib import cm
from matplotlib.animation import FuncAnimation
from matplotlib.colors import LightSource, Normalize, LinearSegmentedColormap

import tkWavefunction_H as wf
import tklib.tkgraphic.tkPlot3d as plt3d


a0 = physical_constants['Bohr radius'][0] * 1e10  # オングストローム単位に変換
pi = np.pi


mode = 'rplot'
functype = '2px'

# for Monte Carlo sampling
nsamples = 30000

figsize = (12, 10)
fontsize = 16


nbins = 100
colors_wf    = [(0, 0, 1), (0, 1, 0), (1, 0, 0)]  # 緑、青、赤
colors_phase = [(1, 0, 0), (0, 1, 0), (0, 0, 1), (1, 1, 0), (1, 0, 0)]  # 赤, 緑, 青, 黄, 赤
cmap_name_wf    = 'custom cmap for wavefunction'
cmap_name_phase = 'custom cmap for phase'


class CustomArgumentParser(argparse.ArgumentParser):
    """
    `argparse.ArgumentParser` を拡張し、カスタムのヘルプメッセージを提供するクラス。

    概要:
    標準のヘルプ出力に加えて、一般的な使用方法と具体的な使用例を追記します。
    """
    def print_help(self, *args, **kwargs):
        """
        カスタムのヘルプメッセージを出力します。

        詳細説明:
        親クラスの `print_help` メソッドの実行前後に、
        本スクリプトの一般的な使用法と具体的な使用例をコンソールに表示します。

        Parameters:
            *args: `argparse.ArgumentParser.print_help` に渡される可変長引数。
            **kwargs: `argparse.ArgumentParser.print_help` に渡されるキーワード引数。

        Returns:
            None
        """
        print()
        print("General usage:")
        super().print_help(*args, **kwargs)
        print()
        print("Examples:")
        print(r"usage: python .\plot_wf.py [rplot|ranim|dot|dotc|dotanim|iso3d|iso3danim|rsurface] [100r|2pxa etc]")
        print(r"usage: python .\plot_wf.py [Rnl|Ylm|Theta|Phi|normalize]")
        print(r"usage: python .\plot_wf.py [3dMC_animation|contour3d|torus]")

# argparse のセットアップ
parser = CustomArgumentParser(description='Plot wave functions of H\n\n')
parser.add_argument('mode', nargs='?', default=mode, help='Mode of operation')
parser.add_argument('functype', nargs='?', default=functype, help='Function type')
parser.add_argument('nsamples', nargs='?', type=int, default=nsamples, help='Number of samples for Monte Carlo sampling')
parser.add_argument('--fplot', type=int, default=1, help='Flag to plot graphs (1: enable, 0: disable)')
parser.add_argument('--fsave', type=int, default=1, help='Flag to save figure file (1: enable, 0: disable)')

# 引数の解析
args = parser.parse_args()

# 引数の設定
mode = args.mode
functype = args.functype
nsamples = args.nsamples
fplot = args.fplot
fsave = args.fsave

# 使用例の出力
print(f"mode = {mode}")
print(f"functype = {functype}")
print(f"nsamples = {nsamples}")
print(f"fplot = {fplot}")


def usage():
    """
    コマンドライン引数のヘルプメッセージを表示します。

    詳細説明:
    `CustomArgumentParser` クラスの `print_help` メソッドを呼び出し、
    プログラムの使用法と例をコンソールに出力します。

    Returns:
        None
    """
    parser.print_help()

def torus(theta: np.ndarray, phi: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """
    トーラスの3D座標を計算します。

    詳細説明:
    トーラスの大半径Rと小半径rを用いて、与えられた角度からX, Y, Z座標を計算します。

    Parameters:
        :param theta: トーラスの縦方向の角度 (0から2π) を表すNumPy配列。
        :param phi: トーラスの横方向の角度 (0から2π) を表すNumPy配列。

    Returns:
        :returns: X, Y, Z座標のタプル。各要素はNumPy配列。
    """
# トーラスのパラメータ
    R = 3  # 大半径
    r = 1  # 小半径

    X = (R + r * np.cos(theta)) * np.cos(phi)
    Y = (R + r * np.cos(theta)) * np.sin(phi)
    Z = r * np.sin(theta)
    return X, Y, Z

def sphere(X: np.ndarray, Y: np.ndarray, Z: np.ndarray) -> np.ndarray:
    """
    3次元空間における点 (X, Y, Z) の半径の二乗を計算します。

    詳細説明:
    各座標成分の二乗和 `X^2 + Y^2 + Z^2` を計算し、
    単位球や一般的な球の等値面を定義する際に利用できます。

    Parameters:
        :param X: X座標のNumPy配列。
        :param Y: Y座標のNumPy配列。
        :param Z: Z座標のNumPy配列。

    Returns:
        :returns: 各点における `X^2 + Y^2 + Z^2` の値。NumPy配列。
    """
    return (X**2 + Y**2 + Z**2)  # 単位球: x^2 + y^2 + z^2 = r^2

def exp_r2(X: np.ndarray, Y: np.ndarray, Z: np.ndarray) -> np.ndarray:
    """
    `exp(-r^2)` の値を計算します。

    詳細説明:
    3次元空間における点 (X, Y, Z) から原点までの距離rの二乗を用いて、
    ガウス型の減衰関数 `exp(-(X^2 + Y^2 + Z^2))` を計算します。

    Parameters:
        :param X: X座標のNumPy配列。
        :param Y: Y座標のNumPy配列。
        :param Z: Z座標のNumPy配列。

    Returns:
        :returns: 各点における `exp(-(X^2 + Y^2 + Z^2))` の値。NumPy配列。
    """
    r2 = X**2 + Y**2 + Z**2
    return np.exp(-r2)


print()
print(f"mode: {mode}")
print(f"function type: {functype}")
ftype, nlmt = wf.analyze_function(functype)
nlm, orb_type = nlmt[:3], nlmt[3]
print(f"   type: {ftype}  nlm:", nlm, f"  orb_type: {orb_type}")

if mode == 'torus':
    f = torus
elif ftype == 'c' or ftype == 'r':
    f = wf.psi_r
else:
    if ftype == 'f':
        if nlm == 'sphere':
            f = sphere
        elif nlm == 'exp_r2':
            f = exp_r2
    elif mode == 'contour3d' or mode == 'iso3d':
        f = sphere
    else:
        f = sphere


def plot_torus() -> matplotlib.pyplot:
    """
    トーラスの3D表面プロットを作成します。

    詳細説明:
    `torus` 関数を用いて座標を生成し、カスタムカラーマップで色付けされた
    トーラスの表面をMatplotlibの3Dプロットで描画します。

    Returns:
        :returns: プロットが描画された `matplotlib.pyplot` モジュール。
    """
# パラメトリック変数
    theta = np.linspace(0, 2 * np.pi, 100)
    phi = np.linspace(0, 2 * np.pi, 100)
    theta, phi = np.meshgrid(theta, phi)

    X, Y, Z = torus(theta, phi)

# カスタムカラーマップの定義
    colors = [(1, 0.5, 0.5), (0, 1, 0), (1, 0.5, 0.5)]  # 薄い赤から青、符水赤へのグラデーション
    color_map = plt3d.make_color_map(phi, colors, cmap_name = 'custom_cmap', nbins = nbins)

# プロット
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')

    plt3d.plot_surface3d(ax, X, Y, Z, facecolors = color_map, edgecolor = 'none', alpha = 0.7, shade = True)

    return plt

def plot_isosurface3d() -> matplotlib.pyplot:
    """
    定義された関数に基づき、複数の等値面を3Dでプロットします。

    詳細説明:
    3Dグリッド上でスカラー場 `f` (グローバル変数) を計算し、
    指定されたレベルで等値面を抽出し、Matplotlibの3Dプロットで描画します。

    Returns:
        :returns: プロットが描画された `matplotlib.pyplot` モジュール。
    """
# 3次元グリッドを作成
    x = np.linspace(-1.2, 1.2, 40)
    y = np.linspace(-1.2, 1.2, 40)
    z = np.linspace(-1.2, 1.2, 40)
    X, Y, Z = np.meshgrid(x, y, z, indexing='ij')

# 3Dスカラー場: 球の等値面 (中心を原点にする)
    F = f(X, Y, Z)

# 等値面のレベル
    levels = [0.1, 0.2, 0.3]
    colors = ['cyan', 'magenta', 'yellow']  # 等値面の色

# 3D プロット
    fig = plt.figure(figsize=(8, 8))
    ax = fig.add_subplot(111, projection='3d')

    dx, dy, dz = x[1] - x[0], y[1] - y[0], z[1] - z[0]
    plt3d.plot_isosurface3d(ax, X, Y, Z, F, levels = levels, origin = [x.min(), y.min(), z.min()], spacing = [dx, dy, dz],
            colors = colors, edgecolor = 'k', alpha = 0.3, linewidth = 0.1)

    return plt

def plot_contour3d() -> matplotlib.pyplot:
    """
    `exp(-r^2)` 関数のスライスされた等高線を3Dでプロットします。

    詳細説明:
    3Dグリッド上で `exp_r2` 関数を計算し、特定のZ座標でのスライスにおける
    等高線 (`F[:, :, 50]`) をMatplotlibの3Dプロットに描画します。

    Returns:
        :returns: プロットが描画された `matplotlib.pyplot` モジュール。
    """
    x = np.linspace(-5, 5, 100)
    y = np.linspace(-5, 5, 100)
    z = np.linspace(-5, 5, 100)
    X, Y, Z = np.meshgrid(x, y, z)
    F = exp_r2(X, Y, Z)

    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')

    plt3d.contour3d(ax, X[:, :, 50], Y[:, :, 50], F[:, :, 50], nbins = nbins, cmap = cm.viridis)

    return plt

def plot_3dMC_animation() -> None:
    """
    モンテカルロサンプリングを用いた3Dアニメーションをプロットします。

    詳細説明:
    確率分布 `|x| * 4π r^3 exp(-r**2)` に基づいてランダムな点を3D空間に順次追加し、
    その分布の変化をアニメーションとして表示します。
    点のX座標の符号に基づいて赤または青に色分けされます。
    アニメーションはユーザーがEnterキーを押すまで継続します。

    Returns:
        :returns: None
    """
# サンプル点の総数と一度に追加する点の数
    num_points = 100000
    nadd = 100

# アニメーションの設定
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')
    sc = ax.scatter([], [], [], marker='o', s=1, alpha=0.5)

# 初期化関数
    def init():
        sc._offsets3d = ([], [], [])
        return sc,

# 更新関数
    def update(frame):
        new_points = nadd  # 追加する新しい点の数
        u = np.random.rand(new_points)
        r = np.cbrt(-np.log(1 - u))  # 立方根で変換

        phi = np.random.rand(new_points) * 2.0 * np.pi
        cos_theta = np.random.rand(new_points) * 2 - 1  # cos(θ) を一様分布させる
        theta = np.arccos(cos_theta)

        x = r * np.sin(theta) * np.cos(phi)
        y = r * np.sin(theta) * np.sin(phi)
        z = r * np.cos(theta)

    # x, y, z の範囲を -10 から 10 にスケーリング
        x = (x / np.max(np.abs(x))) * 10
        y = (y / np.max(np.abs(y))) * 10
        z = (z / np.max(np.abs(z))) * 10

    # 分布関数 |x| * 4π r^3 exp(-r**2)
        weights = np.abs(x) * 4 * np.pi * r**3 * np.exp(-r**2)
        weights /= np.sum(weights)  # 正規化

        indices = np.random.choice(np.arange(new_points), size=new_points, p=weights)
        x_sampled = x[indices]
        y_sampled = y[indices]
        z_sampled = z[indices]

        offsets = np.array(sc._offsets3d)
        offsets = np.vstack([offsets.T, np.column_stack((x_sampled, y_sampled, z_sampled))]).T
        sc._offsets3d = (offsets[0], offsets[1], offsets[2])

    # 色の設定
        color_map = np.where(offsets[0] > 0, 'red', 'blue')  # x > 0 を赤、x < 0 を青
        sc.set_color(color_map)

        ax.set_xlabel('X')
        ax.set_ylabel('Y')
        ax.set_zlabel('Z')

    # 軸の範囲を設定
        ax.set_xlim([-10, 10])
        ax.set_ylim([-10, 10])
        ax.set_zlim([-10, 10])

    # タイトルに点の総数を表示
        ax.set_title(f'Total points: {len(offsets[0])}')

        return sc,

# アニメーションの作成
    ani = FuncAnimation(fig, update, frames=num_points // nadd, init_func=init, blit=False, repeat=False)

    if fplot: plt.pause(0.1)

#ani.save('3dMC_animation.gif', writer='imagemagick')

    input("Press ENTER to terminate>>")

    return None

def plot_Phi() -> matplotlib.pyplot:
    """
    球面調和関数の角度依存部分 `exp(j*m*phi)` の実部と虚部をプロットします。

    詳細説明:
    さまざまな磁気量子数 `m` (0から4) に対して、角度 `phi` に応じた
    `exp(j*m*phi)` の実部と虚部の変化を2つのサブプロットに分けてグラフ化します。

    Returns:
        :returns: プロットが描画された `matplotlib.pyplot` モジュール。
    """
    N = 101
    phi0, phi1 = 0.0, 2.0 * pi
    phi = np.linspace(phi0, phi1, N)

    flist = []
    for m in range(0, 5):
        f = wf.Ylm_phi(m, phi)
        flist.append([m, f.real, f.imag])

    fig, axes = plt.subplots(2, 1, figsize = (10, 10))
 
    phi_deg = 180.0 / pi * phi
    for inf in flist:
        axes[0].plot(phi_deg, inf[1], label = f"m={inf[0]},real")
        axes[1].plot(phi_deg, inf[2], label = f"m={inf[0]},imag")
    axes[0].set_xlabel('phi')
    axes[0].set_ylabel('exp(j*m*phi),real')
    axes[0].legend()

    axes[1].set_xlabel('phi')
    axes[1].set_ylabel('exp(j*m*phi),imag')
    axes[1].legend()
    
    return plt


def plot_Ylm() -> matplotlib.pyplot:
    """
    球面調和関数の角度依存部分 `Ylm_theta(l, m, theta)` をプロットします。

    詳細説明:
    さまざまな方位量子数 `l` (0から2) と磁気量子数 `m` に対して、
    角度 `theta` に応じた `Ylm_theta` の変化をグラフに描画します。

    Returns:
        :returns: プロットが描画された `matplotlib.pyplot` モジュール。
    """
    N = 101
    theta0, theta1 = 0.0, pi
    theta = np.linspace(theta0, theta1, N)

    flist = []
    for l in range(0, 3):
        for m in range(-l, l + 1):
            f = wf.Ylm_theta(l, m, theta)
            flist.append([l, m, f])

    fig, ax = plt.subplots(1, 1, figsize = (10, 10))
 
    theta_deg = 180.0 / pi * theta
    for inf in flist:
        ax.plot(theta_deg, inf[2], label = f"l,m={inf[0]},{inf[1]}")
    ax.set_title('Spherical harmonic (theta dependence)')
    ax.set_xlabel('theta')
    ax.set_ylabel('Ylm')

    ax.legend()
    
    return plt

def plot_Rnl() -> matplotlib.pyplot:
    """
    水素原子の動径波動関数 `Rnl(r)` と動径確率密度分布 `Pnl(r)` をプロットします。

    詳細説明:
    主量子数 `n` と方位量子数 `l` の異なる組み合わせ (1s, 2s, 2p, ..., 4f) について、
    動径波動関数 `Rnl(r)` と `r^2 * |Rnl(r)|^2` (動径確率密度分布) を計算し、
    それぞれを9つのサブプロットを持つ2つの図に描画します。
    `fplot` および `fsave` 引数に応じて表示または保存されます。

    Returns:
        :returns: プロットが描画された `matplotlib.pyplot` モジュール。
    """
    Nr = 101
    rmin = 0.0
    rmax = 40.0

    a00 = 1.0
    r = np.linspace(rmin, rmax, Nr)
    f1list = []
    f2list = []
    P1list = []
    P2list = []
    for inf in [['1s', 1, 0], ['2s', 2, 0], ['2p', 2, 1], ['3s', 3, 0], ['3p', 3, 1], ['3d', 3, 2], 
                ['4s', 4, 0], ['4p', 4, 1], ['4d', 4, 2], ['4f', 4, 3]]:
        f1 = wf.Rnl(r, inf[1], inf[2])
        f2 = wf.Rnl_if(r, inf[1], inf[2])
        f1list.append([inf[0], inf[1], inf[2], f1])
        f2list.append([inf[0], inf[1], inf[2], f2])
        if f1 is not None:
            P1list.append([inf[0], inf[1], inf[2], r**2 * f1**2])
        else:
            P1list.append([inf[0], inf[1], inf[2], None])
        if f2 is not None:
            P2list.append([inf[0], inf[1], inf[2], r**2 * f2**2])
        else:
            P2list.append([inf[0], inf[1], inf[2], None])

#plot
#   plt.rcParams["font.size"] = 18
    fig1, ax1 = plt.subplots(3, 3, figsize = (10, 10))
    ax1 = ax1.flatten()
    fig1.subplots_adjust(hspace=0.4, wspace=0.3)
    fig1.suptitle('$R_{nl}(r)$: Radial wave function')

    for ax, inf1, inf2 in zip(ax1, f1list, f2list):
        if inf1[3] is not None:
            ax.plot(r, inf1[3], label='Module')
        if inf2[3] is not None:
            ax.plot(r, inf2[3], label='explicit', linestyle = 'None', marker='o', markersize = 2.0)
        ax.set_title(inf1[0])
        ax.grid()
        ax.legend()
        ax.set_xlabel('r/$a_0$')
        ax.set_ylabel('$R_{nl}(r)$')

    if fsave: plt.savefig('plot_wf_Rnl.png')

#動径確率密度分布の描画
    fig2, ax2 = plt.subplots(3, 3, figsize = (10, 10))
    ax2 = ax2.flatten()
    fig2.subplots_adjust(hspace=0.4, wspace=0.3)
    fig2.suptitle('$P_{nl}(r)$: Radial probability density distribution')

    for ax, inf1, inf2 in zip(ax2, P1list, P2list):
        if inf1[3] is not None:
            ax.plot(r, inf1[3], label='Module')
        if inf2[3] is not None:
            ax.plot(r, inf2[3], label='explicit', linestyle = 'None', marker='o', markersize = 2.0)
        ax.set_title(inf1[0])
        ax.grid()
        ax.legend()
        ax.set_xlabel('r/$a_0$')
        ax.set_ylabel('$P_{nl}(r)$')

    if fsave: plt.savefig('plot_wf_Pnl.png')

    return plt

def plot_rplot() -> matplotlib.pyplot:
    """
    指定された水素原子の波動関数 (球面調和関数) の3D表面プロットを作成します。

    詳細説明:
    グローバル変数 `nlm` で指定された量子数 `n, l, m` と `orb_type` に基づいて
    波動関数を計算し、その実部、虚部、絶対値、または絶対値の2乗を3D空間にプロットします。
    プロットの色は波動関数の値または位相に基づいてカスタムカラーマップで色分けされます。
    XY, YZ, ZX平面に波動関数の2D等高線も重ねて描画されます。
    波動関数のエネルギー `En` も表示されます。

    Returns:
        :returns: プロットが描画された `matplotlib.pyplot` モジュール。
    """
    n, l, m = nlm
    orb_name = wf.get_orb_name(n, l, m)
    n, l, m = wf.get_qnumbers(n, l, m)

    print()
    print(f"Plot R(theta, phi) for {orb_name}  n, l, m=({n}, {l}, {m})")
    print(f"  orbital type: {orb_type}")
    En = wf.En(n)
    print(f"  En={En:.6g} eV")

    theta = np.linspace(0, pi, 101)
    phi   = np.linspace(0, 2.0 * pi, 101)
    Theta2d, Phi2d = np.meshgrid(theta, phi)
#    r  = wf.Ylm_real(Theta2d, Phi2d, l, m)
    f = wf.Ylm(Theta2d, Phi2d, l, m)
    r, phase = wf.get_by_type(f, orb_type)
    nfigures = 2
    if orb_type == 'a':
        title = r"|$\psi$|"
    elif orb_type == 'a2':
        title = r"|$\psi$|$^2$"
    elif orb_type == 'r':
        title = r"$\psi_r$"
    elif orb_type == 'i':
        title = r"$\psi_i$"

    ra = abs(r)
    z   = ra * cos(Theta2d)
    rxy = ra * sin(Theta2d)
    x   = rxy * cos(Phi2d)
    y   = rxy * sin(Phi2d)

    alpha = 0.7

# カスタムカラーマップの作成
    color_list = colors_wf
    cmap_name = cmap_name_wf
    custom_cmap = LinearSegmentedColormap.from_list(cmap_name, color_list, N = nbins)
# rの値にcmapを対応させる
    norm = Normalize(vmin = min(r), vmax = max(r))
    face_colors = custom_cmap(norm(r))

#    ax = fig.add_subplot(111, projection='3d')
    fig, axes = plt.subplots(1, nfigures, figsize = (8 * nfigures, 10), subplot_kw = {'projection': '3d'})
    if nfigures == 1:
        ax = axes
    else:
        ax = axes[0]
        ax2 = axes[1]

    if phase is None:
        options = {"facecolors": face_colors, "edgecolor": None}
    else:
        phase_min = np.min(phase)
        phase_max = np.max(phase)
        if phase_min == phase_max:
            phase_min -= 0.1
            phase_max += 0.1
        normalized_phase = (phase - phase_min) / (phase_max - phase_min)

# カスタムカラーマップを定義
        colors = colors_phase
        cmap_name = cmap_name_phase
        custom_cmap = LinearSegmentedColormap.from_list(cmap_name, colors, N = nbins)
        facecolors = custom_cmap(normalized_phase)

        options = {"facecolors": facecolors, "edgecolor": None}

    plt3d.plot_surface3d(ax, x, y, z, **options, alpha = alpha, shade = False,
                    xlabel = 'x', ylabel = 'y', zlabel = 'z')

    def contour_func(x, y, z):
        f = wf.Ylm_xyz(x, y, z, l, m)
        r, phase = wf.get_by_type(f, orb_type)
        return r

    maxx = plt3d.get_max_xyz(x, y, z)
    plt3d.plot_contours_xyz_by_func(ax, contour_func, minx = -maxx, maxx = maxx, nmesh = 100,
                            posx = 0.1, posy = 0.1, posz = 0.1,
                            offsetx = -maxx, offsety = maxx, offsetz = -maxx, 
                            cmap = custom_cmap, levels = 20, alpha = 0.3)

    plt3d.set_cubic_scale(ax, maxx)
    ax.set_title(title)

    if phase is None:
        cbar = plt3d.show_color_bar(fig, ax, scale = z, cmap = custom_cmap, shrink = 0.5, aspect = 10.0)
        if orb_type == 'i':
            cbar.set_label(r'$\psi_i$')
        else:
            cbar.set_label(r'$\psi_r$')
    else:
        vmin = phase_min * 180.0 / pi
        vmax = phase_max * 180.0 / pi
        scmap = plt.cm.ScalarMappable(cmap = custom_cmap, norm = plt.Normalize(vmin = vmin, vmax = vmax))
        scmap.set_array(normalized_phase)
        cbar = plt.colorbar(scmap, ax=ax, shrink=0.5, aspect=10.0)
        cbar.set_label('phase')

    if phase is not None and nfigures > 1:
        ra = abs(phase)
        z   = ra * cos(Theta2d)
        rxy = ra * sin(Theta2d)
        x   = rxy * cos(Phi2d)
        y   = rxy * sin(Phi2d)
        plt3d.plot_surface3d(ax2, x * 180.0/pi, y * 180.0/pi, z * 180.0/pi, **options, alpha = alpha, shade = False,
                    xlabel = 'x', ylabel = 'y', zlabel = 'z')
#        plt3d.set_cubic_scale(ax2, maxx)
        ax2.set_title("phase")

    return plt

def plot_ranim() -> None:
    """
    水素原子の波動関数の3D表面プロットをアニメーションとして表示します。

    詳細説明:
    `plot_rplot` と同様に、グローバル変数 `nlm` と `orb_type` に基づいて
    波動関数を計算し、時間依存性 (`phase_m` パラメータ) を加えてアニメーション化します。
    波動関数の位相変化を表現し、GIFファイルとして保存することも可能です。
    アニメーションはユーザーがEnterキーを押すまで継続します。

    Returns:
        :returns: None
    """
    n, l, m = nlm
    orb_name = wf.get_orb_name(n, l, m)
    n, l, m = wf.get_qnumbers(n, l, m)

    print()
    print(f"Animation of R(theta, phi)=Ylm(theta, phi) for {orb_name}  n, l, m=({n}, {l}, {m})")
    print(f"  orbital type: {orb_type}")
    En = wf.En(n)
    print(f"  En={En:.6g} eV")

    theta = np.linspace(0, pi, 101)
    phi   = np.linspace(0, 2.0 * pi, 101)
    Theta2d, Phi2d = np.meshgrid(theta, phi)
#    r  = wf.Ylm_real(Theta2d, Phi2d, l, m)
    f = wf.Ylm(Theta2d, Phi2d, l, m)
    r, phase = wf.get_by_type(f, orb_type)
    nfigures = 1
    if orb_type == 'a':
        title = rf"|$\psi$| ({n}{l}{m})"
    elif orb_type == 'a2':
        title = rf"|$\psi$|$^2$ ({n}{l}{m})"
    elif orb_type == 'r':
        title = rf"$\psi_r$ ({n}{l}{m})"
    elif orb_type == 'i':
        title = rf"$\psi_i$ ({n}{l}{m})"

    ra = abs(r)
    z   = ra * cos(Theta2d)
    rxy = ra * sin(Theta2d)
    x   = rxy * cos(Phi2d)
    y   = rxy * sin(Phi2d)

    alpha = 0.7

    fig, ax = plt.subplots(1, nfigures, figsize = (8, 8), subplot_kw = {'projection': '3d'})

    phase_min = -pi
    phase_max =  pi
    normalized_phase = (phase - phase_min) / (phase_max - phase_min)

# カスタムカラーマップを定義
    colors = colors_phase
    cmap_name = cmap_name_phase
    custom_cmap = LinearSegmentedColormap.from_list(cmap_name, colors, N = nbins)
    facecolors = custom_cmap(normalized_phase)

    options = {"facecolors": facecolors, "edgecolor": None}

    plt3d.plot_surface3d(ax, x, y, z, **options, alpha = alpha, shade = False,
                    xlabel = 'x', ylabel = 'y', zlabel = 'z')

    maxx = plt3d.get_max_xyz(x, y, z)
    plt3d.set_cubic_scale(ax, maxx)
    ax.set_title(title)

    vmin = -180.0
    vmax =  180.0
    scmap = plt.cm.ScalarMappable(cmap = custom_cmap, norm = plt.Normalize(vmin = vmin, vmax = vmax))
    scmap.set_array(normalized_phase)
    cbar = plt.colorbar(scmap, ax = ax, shrink = 0.5, aspect = 10.0)
    cbar.set_label('phase')

    nmaxiter = 100
    wait = 5 # ms
    ani = None
    window_closed = False
    def update_animation(frame):
        nonlocal window_closed, ani

        if window_closed :
            if ani.event_source: ani.event_source.stop()
            return ax

        if frame % 10 == 0: print(f"frame: {frame}/{nmaxiter}")

        ax.cla()  # 軸の設定を保持するために全体をクリア

        f = wf.Ylm(Theta2d, Phi2d, l, m, phase_m = frame * 0.1)
        r, phase = wf.get_by_type(f, orb_type)
        ra = abs(r)
        z   = ra * cos(Theta2d)
        rxy = ra * sin(Theta2d)
        x   = rxy * cos(Phi2d)
        y   = rxy * sin(Phi2d)

        phase_min = -pi
        phase_max =  pi
        normalized_phase = (phase - phase_min) / (phase_max - phase_min)
        custom_cmap = LinearSegmentedColormap.from_list(cmap_name, colors, N = nbins)
        facecolors = custom_cmap(normalized_phase)
        options = {"facecolors": facecolors, "edgecolor": None}

        plt3d.plot_surface3d(ax, x, y, z, **options, alpha = alpha, shade = False,
                    xlabel = 'x', ylabel = 'y', zlabel = 'z')

        plt3d.set_cubic_scale(ax, maxx)
        ax.set_title(f"{title} frame: {frame}/{nmaxiter}")

        if fplot: plt.pause(0.1)

#        if frame >= nmaxiter - 1: 
#            ani.event_source.stop()
#            return ax
        
        return ax

    def on_close(event):
        nonlocal window_closed

        window_closed = True

    fig.canvas.mpl_connect('close_event', on_close)

    ani = FuncAnimation(fig, update_animation, frames = nmaxiter, interval = wait)
#ani = FuncAnimation(fig, update, frames=nmaxiter, init_func=init, blit=False, repeat=False)

    outfile = f'plot_wf_{mode}_{functype}.gif'
    if fsave:
        print(f"saving to {outfile}")
        ani.save(outfile, writer='imagemagick')

    if fplot:
        input("Press Enter to close the window...")

#    plt.show()

    return None


def plot_rsurface() -> matplotlib.pyplot:
    """
    水素原子の動径波動関数の実部を3D表面としてプロットします。

    詳細説明:
    球面座標系における動径波動関数の実部を計算し、それを3D空間の表面として描画します。
    X軸、Y軸、Z軸はそれぞれ直交座標に対応し、色は波動関数の値に基づいて色分けされます。
    XY, YZ, ZX平面に波動関数の2D等高線も重ねて描画されます。

    Returns:
        :returns: プロットが描画された `matplotlib.pyplot` モジュール。
    """
    n, l, m = nlm
    orb_name = wf.get_orb_name(n, l, m)
    n, l, m = wf.get_qnumbers(n, l, m)

    print()
    print(f"Plot R(theta, phi) for {orb_name}  n, l, m=({n}, {l}, {m})")

    theta = np.linspace(0, pi, 101)
    phi   = np.linspace(0, 2.0 * pi, 101)
    Theta2d, Phi2d = np.meshgrid(theta, phi)
    r  = wf.Ylm_real(Theta2d, Phi2d, l, m)
    ra = abs(r)
    z   = ra * cos(Theta2d)
    rxy = ra * sin(Theta2d)
    x   = rxy * cos(Phi2d)
    y   = rxy * sin(Phi2d)

    alpha = 0.7


# カスタムカラーマップの作成
    color_list = colors_wf
    cmap_name = cmap_name_wf
    custom_cmap = LinearSegmentedColormap.from_list(cmap_name, color_list, N = nbins)
# rの値にcmapを対応させる
    norm = Normalize(vmin = min(r), vmax = max(r))
    face_colors = custom_cmap(norm(r))

    fig = plt.figure()
    ax = fig.add_subplot(111, projection = '3d')

#    options = {"colors": colors}
#    options = {"color": color}
#    options = {"cmap": cmap, "edgecolor": None}
    options = {"facecolors": face_colors, "edgecolor": None}
#    options = {"facecolors": colors, "edgecolor": None}
    plt3d.plot_surface3d(ax, x, y, r, **options, alpha = alpha, shade = False,
                    xlabel = 'x', ylabel = 'y', zlabel = 'z')

    maxx = plt3d.get_max_xyz(x, y, z)
    plt3d.plot_contours_xyz_by_func(ax, lambda x, y, z: wf.Ylm_xyz_real(x, y, z, l, m), minx = -maxx, maxx = maxx, nmesh = 100,
                            posx = 0.1, posy = 0.1, posz = 0.1,
                            offsetx = -maxx, offsety = maxx, offsetz = -maxx, 
                            cmap = custom_cmap, levels = 20, alpha = 0.3)

    plt3d.set_cubic_scale(ax, maxx)

    plt3d.show_color_bar(fig, ax, scale = z, cmap = custom_cmap, shrink = 0.5, aspect = 10.0)

    return plt


def estimate_max_f(func, minx: float = -3, maxx: float = -3, nsamples_estimate: int = 1000) -> float:
    """
    3次元空間で定義された関数の絶対値の最大値を推定します。

    詳細説明:
    指定された範囲内でランダムにサンプリングを行い、
    関数の最大値を概算します。これはリジェクションサンプリングの受け入れ率Mを
    決定する際に使用されます。

    Parameters:
        :param func: 最大値を推定する対象の関数 (3つの引数x, y, zを取る)。
        :param minx: X, Y, Z座標の最小値。
        :param maxx: X, Y, Z座標の最大値。
        :param nsamples_estimate: 最大値推定のためのサンプル数。

    Returns:
        :returns: 推定された関数の最大値。
    """
    x_samples, y_samples, z_samples = np.random.uniform(minx, maxx, (3, nsamples_estimate))
    f_vals = func(x_samples, y_samples, z_samples)
    maxf = np.max(f_vals)

    return maxf

def rejection_sampling(func, func2, M: float | None = None,  minx: float = -3, maxx: float = 3, nsamples: int = 10000, nsamples_estimate: int = 1000) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray, float]:
    """
    リジェクションサンプリングを用いて、指定された関数に従う3D点をサンプリングします。

    詳細説明:
    3次元の関数 `func` (絶対値) と `func2` (値と位相) に基づいて、
    ランダムな点を生成し、そのうち条件を満たす点のみを採択します。
    これにより、波動関数の確率密度に比例する点の分布を得ます。
    `M` が指定されない場合、`estimate_max_f` を用いて自動で推定されます。

    Parameters:
        :param func: サンプリングの確率密度として使用する関数（絶対値）。
        :param func2: 波動関数の値と位相を返す関数。
        :param M: サンプリングの受け入れ率の上限値。Noneの場合、`estimate_max_f`で推定される。
        :param minx: サンプリング範囲の最小X, Y, Z座標。
        :param maxx: サンプリング範囲の最大X, Y, Z座標。
        :param nsamples: 試行するサンプル点の総数。
        :param nsamples_estimate: `M`がNoneの場合に`estimate_max_f`で使われるサンプル数。

    Returns:
        :returns: サンプリングされたX, Y, Z座標、関数値、位相のNumPy配列、
                  および最終的に使用されたM値のタプル。
    """
    if M is None: M = estimate_max_f(lambda x, y, z: abs(func(x, y, z)), minx, maxx, nsamples_estimate)

    x_list = []
    y_list = []
    z_list = []
    f_list = []
    phase_list = []
    for _ in range(nsamples):
        x, y, z = np.random.uniform(minx, maxx, 3)
        f_val, p_val = func2(x, y, z)
        if np.random.uniform(0, M) < abs(f_val):
            x_list.append(x)
            y_list.append(y)
            z_list.append(z)
            f_list.append(f_val)
            phase_list.append(p_val)

    return np.array(x_list), np.array(y_list), np.array(z_list), np.array(f_list), np.array(phase_list), M

def plot_dot(plot_contours: bool = True) -> matplotlib.pyplot:
    """
    リジェクションサンプリングで生成された点を用いて、水素原子の波動関数を3D散布図でプロットします。

    詳細説明:
    グローバル変数 `nlm` と `orb_type` で指定された波動関数の確率密度に比例して
    分布する点を3D空間に描画し、点の色は波動関数の位相によって色分けされます。
    `plot_contours` がTrueの場合、波動関数の値に基づく2D等高線もプロットされます。
    各軸のスケールは立方体になるように調整されます。

    Parameters:
        :param plot_contours: 波動関数の等高線プロットを表示するかどうかのブール値。

    Returns:
        :returns: プロットが描画された `matplotlib.pyplot` モジュール。
    """
    nsamples_estimate = 1000
    n, l, m = nlm

    maxx_list = [None, 3.0, 5.0, 10.0, 15.0, 25.0]
    kmaxx_list = [0.5, 0.7, 1.0, 1.0, 1.0, 1.0]
    maxx = maxx_list[n] * kmaxx_list[l]
    minx = -maxx

    offsetx = -maxx
    offsety = maxx
    offsetz = -maxx
    posx = 0.1
    posy = 0.1
    posz = 0.1
    
    if orb_type == 'r':
        title = r'$\psi_r$'
    elif orb_type == 'i':
        title = r'$\psi_i$'
    elif orb_type == 'a':
        title = r'|$\psi$|'
    elif orb_type == 'a2':
        title = r'|$\psi_i$|$^2$'

    print()
    print(f"Plot wave function with dots: n,l,m=({n},{l},{m})")
    print(f"nsamples={nsamples}")


#    func = lambda x, y, z: wf.psi_xyz_real(x, y, z, n, l, m)
    def func(x, y, z):
        f = wf.psi_xyz(x, y, z, n, l, m)
        f, phase = wf.get_by_type(f, orb_type)
        return f

    def func2(x, y, z):
        f = wf.psi_xyz(x, y, z, n, l, m)
        f, phase = wf.get_by_type(f, orb_type)
        return f, phase

    x_sampled, y_sampled, z_sampled, f_sampled, phase_sampled, M = rejection_sampling(func, func2, None, minx, maxx, nsamples, nsamples_estimate)
    nsampled = f_sampled.size

# 正規化の重み
    weights = np.abs(f_sampled)
    weights /= np.sum(weights)

    print(f"nsampled={nsampled}")
    print("Estimated max func:", M)

    alpha_s = 0.2
    alpha_c = 0.1

# 3Dプロットで色を変える
    fig = plt.figure(figsize = (10, 10))
    ax = fig.add_subplot(111, projection = '3d')

# phaseに基づいて色を設定
    colors = colors_phase
    custom_cmap = LinearSegmentedColormap.from_list(cmap_name_phase, colors, N = nbins)
    norm = Normalize(vmin = -pi, vmax = pi)
    scatter = plt3d.plot_scatter3d(ax, x_sampled, y_sampled, z_sampled, minx, maxx, minx, maxx, minx, maxx, 
                    size = 1.0, c = phase_sampled, cmap = custom_cmap, norm = norm, alpha = alpha_s)

    pad = 0.1
    if True:
        scmap = plt.cm.ScalarMappable(cmap = custom_cmap, norm = norm)
        scmap.set_array(phase_sampled)
        cbar = fig.colorbar(scmap, ax = ax, shrink = 0.5, aspect = 10.0, pad = pad)
        cbar.set_label('phase')

    if plot_contours:
#       maxx = plt3d.get_max_xyz(x1, y1, z1)
# カスタムカラーマップの作成
        color_list = colors_wf
        custom_cmap = LinearSegmentedColormap.from_list(cmap_name_wf, color_list, N = nbins)

        fr1d = f_sampled.flatten()
        minF = min(fr1d)
        maxF = max(fr1d)

# rの値にcmapを対応させる
        norm = Normalize(vmin = minF, vmax = maxF)

        plt3d.plot_contours_xyz_by_func(ax, func, minx = -maxx, maxx = maxx, nmesh = 100,
                            posx = posx, posy = posy, posz = posz,
                            offsetx = -maxx, offsety = maxx, offsetz = -maxx, 
                            cmap = custom_cmap, levels = 20, alpha = alpha_c)

        scmap = plt.cm.ScalarMappable(cmap = custom_cmap, norm = norm)
        scmap.set_array(f_sampled)
        cbar = fig.colorbar(scmap, ax = ax, shrink = 0.5, aspect = 10.0, pad = pad)
        cbar.set_label(title)

    plt3d.set_cubic_scale(ax, maxx)

    plt.title(f"dot plot of {title} for nlm={n}{l}{m}")

    return plt

def plot_dotanim(plot_contours: bool = True) -> matplotlib.pyplot:
    """
    リジェクションサンプリングで生成された点を用いて、水素原子の波動関数の3D散布図アニメーションをプロットします。

    詳細説明:
    `plot_dot` と同様に、グローバル変数 `nlm` と `orb_type` に基づいて
    波動関数の確率密度に比例した3D散布図を、時間依存性 (`phase` パラメータ) を加えてアニメーション化します。
    点の色は波動関数の位相によって色分けされ、GIFファイルとして保存することも可能です。
    アニメーションはユーザーがEnterキーを押すまで継続します。
    `plot_contours` 引数は現在この関数では使用されていません。

    Parameters:
        :param plot_contours: 等高線プロットを表示するかどうかのブール値 (アニメーションでは未使用)。

    Returns:
        :returns: プロットが描画された `matplotlib.pyplot` モジュール。
    """
    nsamples_estimate = 1000
    n, l, m = nlm

    maxx_list = [None, 3.0, 5.0, 10.0, 15.0, 25.0]
    kmaxx_list = [0.5, 0.7, 1.0, 1.0, 1.0, 1.0]
    maxx = maxx_list[n] * kmaxx_list[l]
    minx = -maxx

    if orb_type == 'r':
        title = r'$\psi_r$'
    elif orb_type == 'i':
        title = r'$\psi_i$'
    elif orb_type == 'a':
        title = r'|$\psi$|'
    elif orb_type == 'a2':
        title = r'|$\psi_i$|$^2$'

    print()
    print(f"Animation of wave function with dots: n,l,m=({n},{l},{m})")
    print(f"nsamples={nsamples}")

    def func(x, y, z):
        f = wf.psi_xyz(x, y, z, n, l, m)
        f, phase = wf.get_by_type(f, orb_type)
        return f

    def func2(x, y, z):
        f = wf.psi_xyz(x, y, z, n, l, m)
        f, phase = wf.get_by_type(f, orb_type)
        return f, phase

    x_sampled, y_sampled, z_sampled, f_sampled, phase_sampled, M = rejection_sampling(func, func2, None, minx, maxx, nsamples, nsamples_estimate)
    nsampled = f_sampled.size
    print(f"nsampled={nsampled}")
    print("Estimated max func:", M)

# plot
    fig = plt.figure(figsize = (10, 10))
    ax = fig.add_subplot(111, projection = '3d')

# phaseに基づいて色を設定
    colors = colors_phase
    custom_cmap = LinearSegmentedColormap.from_list(cmap_name_phase, colors, N = nbins)
    norm = Normalize(vmin = -pi, vmax = pi)

    scatter = plt3d.plot_scatter3d(ax, x_sampled, y_sampled, z_sampled, minx, maxx, minx, maxx, minx, maxx, 
                    size = 1.0, c = phase_sampled, cmap = custom_cmap, norm = norm)

    scmap = plt.cm.ScalarMappable(cmap = custom_cmap, norm = norm)
    scmap.set_array(phase_sampled)
    cbar = fig.colorbar(scmap, ax = ax, shrink = 0.5, aspect = 10.0)
    cbar.set_label('phase')

    plt3d.set_cubic_scale(ax, maxx)
    plt.title(f"dot plot of {title} for nlm={n}{l}{m}")

    nmaxiter = 62
    wait = 5 # ms
    ani = None
    window_closed = False
    def update_animation(frame):
        nonlocal window_closed, ani

        if window_closed :
            if ani.event_source: ani.event_source.stop()
            return ax

        if frame == 0 and not hasattr(update_animation, 'initialized'):
            update_animation.initialized = True
            return
        if frame % 10 == 0: print(f"frame: {frame}/{nmaxiter}")

        ax.cla()  # 軸の設定を保持するために全体をクリア

        def func2(x, y, z):
            f = wf.psi_xyz(x, y, z, n, l, m, phase = frame * 0.1)
            f, phase = wf.get_by_type(f, orb_type)
            return f, phase

        x_sampled, y_sampled, z_sampled, f_sampled, phase_sampled, _M \
                = rejection_sampling(None, func2, M, minx, maxx, nsamples, nsamples_estimate)

        plt3d.plot_scatter3d(ax, x_sampled, y_sampled, z_sampled, minx, maxx, minx, maxx, minx, maxx, 
                    size = 1.0, c = phase_sampled, cmap = custom_cmap, norm = norm)

        plt3d.set_cubic_scale(ax, maxx)
        ax.set_title(f"{title} frame: {frame}/{nmaxiter}")

        if fplot: plt.pause(0.1)

        return ax

    def on_close(event):
        nonlocal window_closed
        window_closed = True

    fig.canvas.mpl_connect('close_event', on_close)

    ani = FuncAnimation(fig, update_animation, frames = nmaxiter, interval = wait)

    outfile = f'plot_wf_{mode}_{functype}.gif'
    if fsave:
        print(f"saving to {outfile}")
        ani.save(outfile, writer='imagemagick')


    if fplot: 
#    plt.show()
        plt.pause(0.1)
        input("Press Enter to close the window...")

    return plt


def plot_iso3d() -> matplotlib.pyplot:
    """
    水素原子の波動関数の等値面を3Dでプロットします。

    詳細説明:
    グローバル変数 `nlm` と `orb_type` で指定された量子数に基づいて
    3次元グリッド上で水素原子の波動関数を計算し、指定されたレベルで等値面を抽出してプロットします。
    波動関数の値に基づく2D等高線も重ねて描画されます。

    Returns:
        :returns: プロットが描画された `matplotlib.pyplot` モジュール。
    """
    n, l, m = nlm
    orb_name = wf.get_orb_name(n, l, m)
    n, l, m = wf.get_qnumbers(n, l, m)

    maxx_list = [None, 3.0, 5.0, 10.0, 15.0, 25.0]
    kmaxx_list = [0.5, 0.7, 1.0, 1.0, 1.0, 1.0]
    maxx = maxx_list[n] * kmaxx_list[l]

    x = np.linspace(-maxx, maxx, 50)
    y = np.linspace(-maxx, maxx, 50)
    z = np.linspace(-maxx, maxx, 50)
    X, Y, Z = np.meshgrid(x, y, z)

    def func(x, y, z):
        f = wf.psi_xyz(x, y, z, n, l, m)
        f, phase = wf.get_by_type(f, orb_type)
        return f, phase

    def func1(x, y, z):
        return func(x, y, z)[0]

    F, phase = func(X, Y, Z)

# 等値面のレベル
    fr1d = F.flatten()
    minF = min(fr1d)
    maxF = max(fr1d)
    print(f"F range: {minF} - {maxF}")
    levels = [-0.6 * maxF, -0.3 * maxF, 0.3 * maxF, 0.6 * maxF]
    colors = ['blue', 'cyan', 'pink', 'red']  # 等値面の色

# 3D プロット
    fig = plt.figure(figsize = (8, 8))
    ax = fig.add_subplot(111, projection='3d')

    dx, dy, dz = x[1] - x[0], y[1] - y[0], z[1] - z[0]
    plt3d.plot_isosurface3d(ax, X, Y, Z, F, levels = levels, origin = [x.min(), y.min(), z.min()], spacing = [dx, dy, dz],
            colors = colors, edgecolor = 'k', alpha = 0.3, linewidth = 0.1)

# 2D等高線
# カスタムカラーマップの作成
    color_list = colors_wf
    cmap_name = cmap_name_wf
    custom_cmap = LinearSegmentedColormap.from_list(cmap_name, color_list, N = nbins)
# rの値にcmapを対応させる
#    norm = Normalize(vmin = -maxF, vmax = maxF)
    norm = Normalize(vmin = minF, vmax = maxF)
    face_colors = custom_cmap(norm(F))

    maxx = plt3d.get_max_xyz(X, Y, Z)
    plt3d.plot_contours_xyz_by_func(ax, func1, minx = -maxx, maxx = maxx, nmesh = 100,
                            posx = 0.1, posy = 0.1, posz = 0.1,
                            offsetx = -maxx, offsety = maxx, offsetz = -maxx, 
                            cmap = custom_cmap, levels = 20, alpha = 0.3)

    plt3d.set_cubic_scale(ax, maxx)

    scmap = plt.cm.ScalarMappable(cmap = custom_cmap, norm = norm)
    scmap.set_array(F)
    cbar = plt.colorbar(scmap, ax=ax, shrink=0.5, aspect=10.0)
    plt.title(f"isosurfaces for nlm={n}{l}{m}")

    return plt

def plot_iso3danim() -> matplotlib.pyplot:
    """
    水素原子の波動関数の等値面アニメーションを3Dでプロットします。

    詳細説明:
    `plot_iso3d` と同様に、グローバル変数 `nlm` と `orb_type` に基づいて
    波動関数の等値面を計算し、時間依存性 (`phase` パラメータ) を加えてアニメーション化します。
    波動関数の位相変化を表現し、GIFファイルとして保存することも可能です。
    アニメーションはユーザーがEnterキーを押すまで継続します。

    Returns:
        :returns: プロットが描画された `matplotlib.pyplot` モジュール。
    """
    n, l, m = nlm
    orb_name = wf.get_orb_name(n, l, m)
    n, l, m = wf.get_qnumbers(n, l, m)

    if orb_type == 'r':
        title = r'$\psi_r$'
    elif orb_type == 'i':
        title = r'$\psi_i$'
    elif orb_type == 'a':
        title = r'|$\psi$|'
    elif orb_type == 'a2':
        title = r'|$\psi_i$|$^2$'

    maxx_list = [None, 3.0, 5.0, 10.0, 15.0, 25.0]
    kmaxx_list = [0.5, 0.7, 1.0, 1.0, 1.0, 1.0]
    maxx = maxx_list[n] * kmaxx_list[l]

    x = np.linspace(-maxx, maxx, 50)
    y = np.linspace(-maxx, maxx, 50)
    z = np.linspace(-maxx, maxx, 50)
    X, Y, Z = np.meshgrid(x, y, z)

    def func(x, y, z):
        f = wf.psi_xyz(x, y, z, n, l, m)
        f, phase = wf.get_by_type(f, orb_type)
        return f, phase

    def func1(x, y, z):
        return func(x, y, z)[0]

    F, phase = func(X, Y, Z)

# 等値面のレベル
    fr1d = F.flatten()
    minF = min(fr1d)
    maxF = max(fr1d)
    print(f"F range: {minF} - {maxF}")
    levels = [-0.6 * maxF, -0.3 * maxF, 0.3 * maxF, 0.6 * maxF]
    colors = ['blue', 'cyan', 'pink', 'red']  # 等値面の色

# 3D プロット
    fig = plt.figure(figsize = (8, 8))
    ax = fig.add_subplot(111, projection='3d')

    dx, dy, dz = x[1] - x[0], y[1] - y[0], z[1] - z[0]
    plt3d.plot_isosurface3d(ax, X, Y, Z, F, levels = levels, origin = [x.min(), y.min(), z.min()], spacing = [dx, dy, dz],
            colors = colors, edgecolor = 'k', alpha = 0.3, linewidth = 0.1)

    plt3d.set_cubic_scale(ax, maxx)
    plt.title(f"isosurfaces for nlm={n}{l}{m}")

    nmaxiter = 62
    wait = 5 # ms
    ani = None
    window_closed = False
    def update_animation(frame):
        nonlocal window_closed, ani

        if window_closed :
            if ani.event_source: ani.event_source.stop()
            return ax

        if frame == 0 and not hasattr(update_animation, 'initialized'):
            update_animation.initialized = True
            return
        if frame % 10 == 0: print(f"frame: {frame}/{nmaxiter}")

        ax.cla()  # 軸の設定を保持するために全体をクリア

        def func(x, y, z):
            f = wf.psi_xyz(x, y, z, n, l, m, phase = frame * 0.1)
            f, phase = wf.get_by_type(f, orb_type)
            return f, phase

        F, phase = func(X, Y, Z)
        plt3d.plot_isosurface3d(ax, X, Y, Z, F, levels = levels, origin = [x.min(), y.min(), z.min()], spacing = [dx, dy, dz],
                colors = colors, edgecolor = 'k', alpha = 0.3, linewidth = 0.1)

        plt3d.set_cubic_scale(ax, maxx)
        ax.set_title(f"{title} frame: {frame}/{nmaxiter}")

        if fplot: plt.pause(0.1)

        return ax

    def on_close(event):
        nonlocal window_closed
        window_closed = True

    fig.canvas.mpl_connect('close_event', on_close)

    ani = FuncAnimation(fig, update_animation, frames = nmaxiter, interval = wait)

#    plt.show()

    outfile = f'plot_wf_{mode}_{functype}.gif'
    if fsave:
        print(f"saving to {outfile}")
        ani.save(outfile, writer='imagemagick')

    if fplot: 
#    plt.show()
#        plt.pause(0.1)
        input("Press Enter to close the window...")

    return plt


def make_3d_density_map(output_file: str = "EDens.vasp", Z: float = 1.0, n: int = 1, l: int = 0, m: int = 0, lx: float = 10.0, ly: float = 10.0, lz: float = 10.0, step: float = 0.2) -> int:
    """
    水素原子の電子密度を3DグリッドデータとしてVASP形式で出力します。

    詳細説明:
    指定された量子数 `n, l, m` と原子番号 `Z` の波動関数に基づいて、
    X, Y, Z方向のグリッド上に電子密度 (波動関数の絶対値の2乗) を計算し、
    VASP (LOCPOT) ファイル形式で出力します。

    Parameters:
        :param output_file: 出力ファイル名。Orbital名が埋め込まれる（例: "EDens_{Orb}.vasp"）。
        :param Z: 原子番号 (陽子数)。
        :param n: 主量子数。
        :param l: 方位量子数。
        :param m: 磁気量子数。
        :param lx: X方向のボックスサイズ。
        :param ly: Y方向のボックスサイズ。
        :param lz: Z方向のボックスサイズ。
        :param step: グリッドのステップサイズ。

    Returns:
        :returns: 成功時は `0`、ファイル書き込みエラー時は `-1`、
                  無効な量子数の場合は `-2` を返します。
    """
    print("\nCalc 3D Electron Density map\n")

    OrbName = wf.GetOrbName(n, l, m)
    print(f"  Z={Z}")
    print(f"  n,l,m={n},{l},{m}")
    print(f"  Orbital: {OrbName}")
    output_file = output_file.replace("{Orb}", OrbName)
    print(f"OutFile: [{output_file}]")

    try:
        out = open(output_file, "w")
    except IOError:
        print(f"Error: Can not write to [{output_file}].")
        return -1

    nx = int(lx / step + 0.1)
    ny = int(ly / step + 0.1)
    nz = int(lz / step + 0.1)
    print(f"Range: {lx} {ly} {lz}")
    print(f"nMesh: {nx} {ny} {nz}")

    out.write(f"Electron density for Z={Z}, n={n}, l={l}, m={m}\n")
    out.write(f"   {lx:12.6f}\n")
    out.write(f"     {lx/lx:12.6f}   {0.0:12.6f}   {0.0:12.6f}\n")
    out.write(f"     {0.0:12.6f}   {ly/lx:12.6f}   {0.0:12.6f}\n")
    out.write(f"     {0.0:12.6f}   {0.0:12.6f}   {lz/lx:12.6f}\n")
    out.write("  1\n")
    out.write("Direct\n")
    out.write("  0.500000  0.500000  0.500000\n")
    out.write("\n")
    out.write(f"  {nx:3d}  {ny:3d}  {nz:3d}\n")

    count = 0
    for ix in range(nx):
        for iy in range(ny):
            for iz in range(nz):
                x = ix * step
                y = iy * step
                z = iz * step
                # 中心は(0.5*lx, 0.5*ly, 0.5*lz)
                f = wf.cal_wave_function(Z, n, l, m, x, y, z, 0.5*lx, 0.5*ly, 0.5*lz)
                if f is None:
                    print(f"Invalid n,l,m({n},{l},{m}).")
                    out.close()
                    return -2

                out.write(f" {f:11.5g}")
                count += 1
                if count % 10 == 0:
                    out.write("\n")

    out.close()

def integ_MC3d(f, x_min: float, x_max: float, y_min: float, y_max: float, z_min: float, z_max: float, n_samples: int) -> float:
    """
    モンテカルロ法（リジェクトサンプリングではない）を用いて、3次元空間の関数を積分します。

    詳細説明:
    指定された関数の3次元積分を、指定された範囲内でランダムにサンプリングした点での
    関数値の平均に体積を掛けることで近似します。

    Parameters:
        :param f: 積分する対象の関数 (3つの引数x, y, zを取る)。
        :param x_min: X軸の最小値。
        :param x_max: X軸の最大値。
        :param y_min: Y軸の最小値。
        :param y_max: Y軸の最大値。
        :param z_min: Z軸の最小値。
        :param z_max: Z軸の最大値。
        :param n_samples: サンプリング点の数。

    Returns:
        :returns: 近似された積分の値。
    """
    x_samples = np.random.uniform(x_min, x_max, n_samples)
    y_samples = np.random.uniform(y_min, y_max, n_samples)
    z_samples = np.random.uniform(z_min, z_max, n_samples)
    
    S = sum(f(x_samples, y_samples, z_samples))
    volume = (x_max - x_min) * (y_max - y_min) * (z_max - z_min)
    integral = S * volume / n_samples

    return integral

def normalize() -> None:
    """
    水素原子の波動関数が正規化されているかを確認します。

    詳細説明:
    動径波動関数 `Rnl(r)` と全波動関数 `psi(x,y,z)` の両方について、
    `|psi|^2` (確率密度) が全空間で1に積分されることを検証します。
    動径部分の積分には `scipy.integrate.quad` を、3次元部分の積分には
    モンテカルロ積分を使用します。異なる量子数 `n, l` の組み合わせについて検証結果を出力します。

    Returns:
        :returns: None
    """
    print("\nCalculate |psi(x,y,z)|^2 to check normalization\n")
    print("Rnl:")
    for n in range(1, 5):
        for l in range(0, n):
            def func(r):
#                return r**2 * abs(wf.psi_r(r, 0.0, 0.0, n, l, 0))**2
                return 4 * pi * r**2 * abs(wf.psi_r(r, 0.0, 0.0, n, l, 0))**2
#                return 4 * pi * r**2 * wf.Rnl(r, n, 0)**2

            result, error = quad(func, 0, np.inf)
            print(f"  n,l={n},{l}: S={result} +- {error}")

    nsampling = 1000000
    print()
    print("psi:")
    for n in range(1, 5):
        lmax = n - 1
        for l in range(0, lmax + 1):
            print(f"Orbital: {n}{l}0: En={wf.En(n):.6g} eV", end = '')

            def psi2(x, y, z):
                f = wf.psi_xyz(x, y, z, n, l, 0)
                return f.real**2 + f.imag**2

            if n == 1:
                maxx = 5.0
            else:
                maxx = 20.0
            integ = integ_MC3d(psi2, -maxx, maxx, -maxx, maxx, -maxx, maxx, nsampling)
            print(f"  |psi|^2={integ}")

def main(mode: str) -> None:
    """
    プログラムの主要な実行ロジックを制御します。

    詳細説明:
    コマンドライン引数から受け取った `mode` に基づいて、
    対応するプロット関数または処理関数 (例: `plot_torus`, `plot_Rnl`, `normalize`) を呼び出します。
    結果のプロット表示 (`fplot`) やファイル保存 (`fsave`) も管理します。
    無効な `mode` が指定された場合はエラーメッセージを出力し終了します。

    Parameters:
        :param mode: 実行モードを指定する文字列。

    Returns:
        :returns: None
    """
    if mode == 'torus':
        plt = plot_torus()
    elif mode == 'Rnl':
        plt = plot_Rnl()
    elif mode == 'Ylm' or mode == 'Theta':
        plt = plot_Ylm()
    elif mode == 'Phi':
        plt = plot_Phi()
    elif mode == 'normalize':
        normalize()
        exit()
    elif mode == 'rplot':
        plt = plot_rplot()
    elif mode == 'ranim':
        plt = plot_ranim()
    elif mode == 'dot':
        plt = plot_dot(plot_contours = False)
    elif mode == 'dotc':
        plt = plot_dot(plot_contours = True)
    elif mode == 'dotanim':
        plt = plot_dotanim(plot_contours = False)
    elif mode == 'iso3d':
        plt = plot_iso3d()
    elif mode == 'iso3danim':
        plt = plot_iso3danim()
    elif mode == 'rsurface':
        plt = plot_rsurface()
    elif mode == 'isosurface3d':
        plt = plot_isosurface3d()
    elif mode == 'contour3d':
        plt = plot_contour3d()
    elif mode == '3dMC_animation':
        plt = plot_3dMC_animation()
    else:
        print(f"\nError: Invalid mode [{mode}]")
        exit()

    if plt:
        if fsave: 
            output_file = f"plot_wf_{mode}_{functype}.png"
            print()
            print(f"Save plot to {output_file}")
            plt.savefig(output_file)
        if fplot: plt.show()

if __name__ == "__main__":
    main(mode)
    usage()


'''
# ヒストグラムの表示
    if 'h' in mode:
        fig, axes = plt.subplots(2, 2)
        axes[0, 0].hist(r_sampled, bins=100, color='blue', alpha=0.7, density=True)
        axes[0, 0].set_xlabel('r')
        axes[0, 0].set_ylabel('Density')
        axes[1, 0].hist(theta, bins=100, color='blue', alpha=0.7, density=True)  # 修正
        axes[1, 0].set_xlabel('θ')
        axes[1, 0].set_ylabel('Density')
        axes[0, 1].hist(phi, bins=100, color='blue', alpha=0.7, density=True)  # 修正
        axes[0, 1].set_xlabel('φ')
        axes[0, 1].set_ylabel('Density')

        plt.tight_layout()
        plt.show()
        exit()
'''