import os
import numpy as np
import multiprocessing
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from matplotlib.animation import FuncAnimation

# パラメータ設定
k = 2 * np.pi  # 波数 (単位: m^-1)
omega = 2 * np.pi  # 角周波数 (単位: rad/s)
c = 1  # 波動の伝播速度 (単位: m/s)
t_max = 10  # 最大時間
t_steps = 200  # 時間ステップ数
x_range = np.linspace(-5, 5, 100)  # 空間の範囲（x軸）
y_range = np.linspace(0, 5, 100)   # 空間の範囲（y軸）y >= 0に制限

# 空間メッシュグリッドの生成
x, y = np.meshgrid(x_range, y_range)

# 光源の位置 (光源1, 光源2)
source1 = np.array([-1, 0])  # 光源1の位置 (-1, 0)
source2 = np.array([1, 0])   # 光源2の位置 (1, 0)

# 波の干渉計算
def interference_wave(x, y, t, k, omega, source1, source2):
    # 光源1からの波
    r1 = np.sqrt((x - source1[0])**2 + (y - source1[1])**2)
    wave1 = np.exp(1j * (k * r1 - omega * t))
    
    # 光源2からの波
    r2 = np.sqrt((x - source2[0])**2 + (y - source2[1])**2)
    wave2 = np.exp(1j * (k * r2 - omega * t))
    
    # 干渉結果
    total_wave = wave1 + wave2
    return total_wave

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

# 光源位置に球を描画
ax.scatter(source1[0], source1[1], 0, color="red", s=100, label="Source 1")  # 光源1
ax.scatter(source2[0], source2[1], 0, color="blue", s=100, label="Source 2")  # 光源2

# 時間ステップにおける波の干渉をプロット
t = 0  # 初期時間
X, Y = np.meshgrid(x_range, y_range)

# アニメーションの作成
def animate(t):
    # 干渉波を計算
    total_wave = interference_wave(X, Y, t, k, omega, source1, source2)

    # 干渉波の強度（|ψ|^2）を計算
    intensity = np.abs(total_wave)**2

    # y = 5の位置での強度を取得
    # y = 5 に対応するインデックスを取得
    y_index = np.abs(y_range - 5).argmin()
    
    # y = 5 の位置での強度を抽出し、2D配列として生成
    Z_contour = np.zeros_like(intensity)
    Z_contour[y_index, :] = intensity[y_index, :]  # y = 5の位置だけ強度を設定

    # y = 5 での等高線を描画
    ax.clear()
    ax.scatter(source1[0], source1[1], 0, color="red", s=100, label="Source 1")  # 光源1
    ax.scatter(source2[0], source2[1], 0, color="blue", s=100, label="Source 2")  # 光源2

    ax.plot_surface(X, Y, np.real(total_wave), cmap='Blues', edgecolor='none')
    ax.contour(X, Y, Z_contour, 20, cmap='hot')  # 等高線の追加
    
    ax.set_xlim([-5, 5])
    ax.set_ylim([0, 5])
    ax.set_zlim([-5, 5])
    
    ax.set_xlabel('X')
    ax.set_ylabel('Y')
    ax.set_zlabel('Wave Amplitude')
    ax.set_title(f'Wave Interference at time t={t:.2f}s')
    ax.legend()

# プロット
from matplotlib.animation import FuncAnimation
ani = FuncAnimation(fig, animate, frames=np.linspace(0, t_max, t_steps), interval=100)

plt.show()