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

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

# Grid数、スナップショット数取得
num_grid = 22
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()
ax = fig.add_subplot(111, projection='3d')

def plot(num):
    # 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("taiko", size = 20)

    # 軸ラベルを設定
    ax.set_xlabel("x", size = 14, color = "r")
    ax.set_ylabel("y", size = 14, color = "r")
    ax.set_zlabel("u", size = 14, color = "r")

    # 軸範囲の設定
    ax.set_xlim(0, 10)
    ax.set_ylim(0, 10)
    ax.set_zlim(-2, 2)

    # frames 出力
    surface = ax.plot_wireframe(X1,X2,U)
    #surface = ax.plot_surface(X1,X2,PHI)
    #fig.colorbar(surface)


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

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

plt.show()