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_tmp, data_num = np.loadtxt("./outmont2D.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_tmp[num*num_grid2:(num+1)*num_grid2], (X1, X2))


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

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

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

    # 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=150)

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

plt.show()