import numpy as np
import matplotlib.pyplot as plt
from scipy.interpolate import griddata
from matplotlib.animation import FuncAnimation
from matplotlib.gridspec import GridSpec
from mpl_toolkits.mplot3d import axes3d, Axes3D

# データ取得
data_x, data_y, data_u, data_psiim, data_pot = np.loadtxt("./output_quantum1_kadai.dat", comments='!', unpack=True)

# Grid数、スナップショット数取得
num_grid = 122
num_grid2 = num_grid*num_grid
num_tot = len(data_x)
num_snapshots = int(num_tot/num_grid2)

print(num_tot)
print(num_snapshots)

# 図、枠（軸）の設定
fig = plt.figure()
gs = GridSpec(1, 2, width_ratios = [0.9, 0.05])
ax = fig.add_subplot(gs[0], projection='3d') # グラフの枠
ax_cb = fig.add_subplot(gs[1]) # colorbar の枠

def plot(num):
    #fig.clear()
    # grid data 生成

    X1,X2 = np.meshgrid(data_y[num*num_grid:(num+1)*num_grid],data_y[num*num_grid:(num+1)*num_grid])
    U = griddata((data_x[num*num_grid2:(num+1)*num_grid2], data_y[num*num_grid2:(num+1)*num_grid2]), data_u[num*num_grid2:(num+1)*num_grid2], (X1, X2))


    # clear previous frames
    ax.cla()
    # Axesのタイトルを設定
    ax.set_title("2D wave fn. scattering", size = 20)

    # 軸ラベルを設定
    ax.set_xlabel("x", size = 14, color = "b")
    ax.set_ylabel("y", size = 14, color = "b")
    ax.set_zlabel("$|\psi|^2$", size = 14, color = "b")

    # 軸範囲の設定
    ax.set_xlim(-1, 11)
    ax.set_ylim(-1, 11)
    ax.set_zlim(-1.1, 1.1)

    # frames 出力
    surface = ax.plot_surface(X1,X2,U,cmap="jet", vmin=-0.2, vmax=1.0)

    # colorbar 設定
    fig.colorbar(surface, shrink= 0.5, aspect = 5, cax=ax_cb)


#動画生成
ani = FuncAnimation(fig, plot, frames=num_snapshots, interval=50)

#動画保存
#ani.save("quantum1_kagai.gif", writer="pillow")
ani.save('quantum1_kadai.mp4', writer="ffmpeg")

plt.show()