import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

# -----------------------------
# パラメータ設定
# -----------------------------
m = 1.0
kx, ky, kz = 1.0, 2.0, 3.0
E = 3.0
P_max = 0.03  # 発散防止のための最大確率密度

# 許容範囲
margin = 0.3
rx_max = np.sqrt(2 * E / (m * kx))
ry_max = np.sqrt(2 * E / (m * ky))
rz_max = np.sqrt(2 * E / (m * kz))

nx, ny, nz = 60, 60, 60
x = np.linspace(-rx_max*(1+margin), rx_max*(1+margin), nx)
y = np.linspace(-ry_max*(1+margin), ry_max*(1+margin), ny)
z = np.linspace(-rz_max*(1+margin), rz_max*(1+margin), nz)
X, Y, Z = np.meshgrid(x, y, z, indexing='ij')

# -----------------------------
# ポテンシャルと確率密度
# -----------------------------
V = 0.5 * m * (kx * X**2 + ky * Y**2 + kz * Z**2)
P = np.zeros_like(V)
mask = V < E
P[mask] = 1.0 / np.sqrt(2 * (E - V[mask]) / m)
P /= np.max(P)  # 正規化
P = np.clip(P, 0, P_max)  # 発散防止のため最大値でクリップ

# -----------------------------
# z=0 スライスの surface plot
# -----------------------------
z_index = nz // 2
X_slice = X[:, :, z_index]
Y_slice = Y[:, :, z_index]
P_slice = P[:, :, z_index]

fig = plt.figure(figsize=(8, 6))
ax = fig.add_subplot(111, projection='3d')
surf = ax.plot_surface(X_slice, Y_slice, P_slice, cmap='Blues', edgecolor='none')
ax.set_xlabel("x")
ax.set_ylabel("y")
ax.set_zlabel("Probability Density")
ax.set_zscale('log')
ax.set_title("3D Surface Plot (z = 0 slice)")
fig.colorbar(surf, shrink=0.6, aspect=10, label="Normalized Probability")
plt.tight_layout()
plt.show()