# interfere_coherence.py
"""
# 必要ライブラリ
pip install numpy matplotlib

# 例1: N=2, コヒーレンス長=100*λ でヤング干渉（ほぼコヒーレント）
python interfere_coherence.py --N 2 --Lc-multiple 100

# 例2: 同じ条件で Lc=2*λ（不完全コヒーレンス）
python interfere_coherence.py --N 2 --Lc-multiple 2

# 例3: N={2,10,50} と Lc/λ={2,10,100} のグリッド比較
python interfere_coherence.py --grid-N 2 10 50 --grid-Lc 2 10 100

# 例4: コヒーレンス長を0.5*λ → 200*λまで連続変化させて見る（簡易アニメ）
python interfere_coherence.py --N 20 --animate --Lc-range 0.5 200 --frames 120

N を増やすほど主極大は鋭く高く、副極大は低くなります（格子因子の基本挙動）。

Lc を短くするほど、遠いスリット同士の位相相関が失われ、干渉縞が洗われて最終的には単スリットのsinc²エンベロープに近づきます。

d/λ を大きくすると極大の間隔が狭く、a/λ を大きくするとエンベロープが狭くなります。

"""


import argparse
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation

def gaussian_degree_of_coherence(dL, Lc):
    """ガウス型の複素相関（実数化）。Lc→∞ で 1 に漸近。"""
    if np.isinf(Lc) or Lc <= 0:
        # Lc<=0 は完全コヒーレント扱い
        return np.ones_like(dL)
    return np.exp(-0.5 * (dL / Lc) ** 2)

def sinc_envelope(theta, lam, a):
    """
    単スリット回折のエンベロープ: sinc^2( (pi*a/lam) * sinθ )
    numpyのnp.sinc(x)は sin(pi x)/(pi x) なので変換に注意。
    """
    x = (np.pi * a / lam) * np.sin(theta)
    # 安定な sinc 実装
    y = np.ones_like(x)
    nz = np.abs(x) > 1e-16
    y[nz] = np.sin(x[nz]) / x[nz]
    return y**2

def intensity_multi_slit(theta, N, d, a, lam, Lc):
    """
    Nスリット（等間隔 d、スリット幅 a、波長 lam）の角度分布 I(θ)。
    有限コヒーレンス長 Lc をガウス相関で掛ける。
    """
    k = 2 * np.pi / lam
    phi = k * d * np.sin(theta)  # 位相差 φ(θ)
    env = sinc_envelope(theta, lam, a)

    # 格子因子（有限コヒーレンス）
    # I = N  +  2 * Σ_{p=1..N-1} (N-p) * γ(p d sinθ) * cos(p φ)
    I_core = np.full_like(theta, fill_value=N, dtype=float)
    # ペア間隔 p ごとに寄与を足す（ベクトル化）
    # 角度ベクトル theta に対し、p ごとに cos(p*phi(theta)) と γ を掛ける
    for p in range(1, N):
        dL = p * d * np.sin(theta)  # 光路差
        gamma = gaussian_degree_of_coherence(dL, Lc)
        I_core += 2.0 * (N - p) * gamma * np.cos(p * phi)

    # 単スリット回折のエンベロープを掛ける
    return env * I_core

def make_theta(max_deg=60.0, n=2001):
    """-max_deg～+max_deg の角度配列（ラジアン）"""
    th = np.linspace(np.deg2rad(-max_deg), np.deg2rad(+max_deg), n)
    return th

def plot_single(N, Lc_multiple, lam, d, a, theta_max_deg=60.0, ax=None, title=None):
    """
    単一条件（N, Lc）で強度 I(θ) を描画。
    Lc_multiple は Lc/λ を与える（例: 100 -> Lc = 100 * lam）。
    """
    th = make_theta(theta_max_deg)
    Lc = np.inf if (Lc_multiple is None or Lc_multiple <= 0) else (Lc_multiple * lam)
    I = intensity_multi_slit(th, N=N, d=d, a=a, lam=lam, Lc=Lc)

    if ax is None:
        fig, ax = plt.subplots(figsize=(8, 4.5))
    ax.plot(np.rad2deg(th), I)
    ax.set_xlabel("Angle θ (deg)")
    ax.set_ylabel("Relative Intensity (arb. unit)")

    if title is None:
        if np.isinf(Lc):
            ctext = "Lc = ∞ (fully coherent)"
        else:
            ctext = f"Lc = {Lc_multiple:.3g} × λ"
        title = f"N = {N}, {ctext}"
    ax.set_title(title)
    ax.grid(True)
    return ax

def plot_grid(N_list, Lc_list, lam, d, a, theta_max_deg=60.0):
    """
    N と Lc/λ の組み合わせグリッドを描画。
    """
    nrows = len(N_list)
    ncols = len(Lc_list)
    fig, axes = plt.subplots(nrows, ncols, figsize=(4.5*ncols, 3.2*nrows), sharex=True, sharey=True)
    if nrows == 1 and ncols == 1:
        axes = np.array([[axes]])
    elif nrows == 1:
        axes = axes[np.newaxis, :]
    elif ncols == 1:
        axes = axes[:, np.newaxis]

    for i, N in enumerate(N_list):
        for j, Lcm in enumerate(Lc_list):
            title = None
            if np.isinf(Lcm) or Lcm <= 0:
                ctext = "Lc = ∞"
            else:
                ctext = f"Lc = {Lcm:g}×λ"
            title = f"N={N}, {ctext}"
            plot_single(N, Lc_multiple=Lcm, lam=lam, d=d, a=a,
                        theta_max_deg=theta_max_deg, ax=axes[i, j], title=title)

    fig.suptitle(f"Multi-slit interference with finite coherence\nλ={lam*1e9:.3g} nm, d={d/lam:.2f} λ, a={a/lam:.2f} λ", y=0.99)
    plt.tight_layout()
    return fig

def animate_Lc(N, lam, d, a, Lc_min_multiple, Lc_max_multiple, frames=100, theta_max_deg=60.0):
    """
    Lc/λ を連続的に変化させる簡易アニメーション。
    """
    th = make_theta(theta_max_deg)
    fig, ax = plt.subplots(figsize=(8, 4.5))
    line, = ax.plot([], [])
    ax.set_xlim(-theta_max_deg, theta_max_deg)
    ax.set_ylim(0, None)
    ax.set_xlabel("Angle θ (deg)")
    ax.set_ylabel("Relative Intensity (arb. unit)")
    ax.grid(True)

    txt = ax.text(0.02, 0.95, "", transform=ax.transAxes, va="top")

    def frame_to_Lc_multiple(f):
        return Lc_min_multiple * (Lc_max_multiple / Lc_min_multiple) ** (f / (frames - 1))

    def init():
        ax.set_title(f"N = {N}  (λ={lam*1e9:.3g} nm, d={d/lam:.2f} λ, a={a/lam:.2f} λ)")
        line.set_data([], [])
        txt.set_text("")
        return line, txt

    def update(f):
        Lcm = frame_to_Lc_multiple(f)
        Lc = np.inf if Lcm <= 0 else Lcm * lam
        I = intensity_multi_slit(th, N=N, d=d, a=a, lam=lam, Lc=Lc)
        line.set_data(np.rad2deg(th), I)
        if np.isinf(Lc):
            label = "Lc = ∞ (fully coherent)"
        else:
            label = f"Lc = {Lcm:.3g}×λ"
        txt.set_text(label)
        return line, txt

    ani = FuncAnimation(fig, update, frames=frames, init_func=init, blit=True, interval=60)
    plt.show()

def main():
    p = argparse.ArgumentParser(description="有限コヒーレンスとスリット数で変わる干渉パターンの可視化")
    # 物理パラメータ
    p.add_argument("--lambda-nm", type=float, default=532.0, help="中心波長 λ [nm]（可視光の例：532）")
    p.add_argument("--d-multiple", type=float, default=3.0, help="スリット間隔 d を λ の何倍にするか（例: 3.0 → d=3λ）")
    p.add_argument("--a-multiple", type=float, default=0.5, help="スリット幅 a を λ の何倍にするか（例: 0.5 → a=0.5λ）")
    p.add_argument("--theta-max", type=float, default=60.0, help="描画角度範囲の上限（±度）")
    # 単一プロット用
    p.add_argument("--N", type=int, default=2, help="スリット数")
    p.add_argument("--Lc-multiple", type=float, default=100.0, help="コヒーレンス長 Lc を λ の何倍にするか（<=0 なら ∞ とみなす）")
    # グリッド比較
    p.add_argument("--grid-N", type=int, nargs="*", help="グリッド表示する N のリスト（例: --grid-N 2 10 50）")
    p.add_argument("--grid-Lc", type=float, nargs="*", help="グリッド表示する Lc/λ のリスト（例: --grid-Lc 2 10 100 0 で 0は∞扱い）")
    # アニメーション
    p.add_argument("--animate", action="store_true", help="Lc/λ を連続的に変えてアニメーション表示")
    p.add_argument("--Lc-range", type=float, nargs=2, metavar=("MIN", "MAX"),
                   default=[0.5, 200.0], help="アニメ時の Lc/λ の範囲（対数的に推移）")
    p.add_argument("--frames", type=int, default=120, help="アニメーションのフレーム数")

    args = p.parse_args()

    lam = args.lambda_nm * 1e-9
    d = args.d_multiple * lam
    a = args.a_multiple * lam

    if args.animate:
        animate_Lc(N=args.N, lam=lam, d=d, a=a,
                   Lc_min_multiple=args.Lc_range[0],
                   Lc_max_multiple=args.Lc_range[1],
                   frames=args.frames, theta_max_deg=args.theta_max)
        return

    if args.grid_N and args.grid_Lc:
        N_list = args.grid_N
        Lc_list = [np.inf if (x <= 0) else x for x in args.grid_Lc]
        plot_grid(N_list, Lc_list, lam=lam, d=d, a=a, theta_max_deg=args.theta_max)
        plt.show()
        return

    # 単一条件
    Lc_multiple = args.Lc_multiple
    plot_single(args.N, Lc_multiple, lam=lam, d=d, a=a, theta_max_deg=args.theta_max)
    plt.show()

if __name__ == "__main__":
    main()