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

# データ取得
data_x, data_tmp = np.loadtxt("./output1D.dat", comments='#', unpack=True)
data_xM, data_tmpM, data_num = np.loadtxt("./outmont1D.dat", comments='#', unpack=True)

# グリッド数、スナップショット数取得
num_grid = 22
num_tot = len(data_x)
num_snapshots = int(num_tot/num_grid)

print(num_tot)
print(num_snapshots)

# 図、枠（軸）の設定
fig = plt.figure()
ax = fig.add_subplot(111)


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

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

    # 軸目盛を設定
#    ax.set_xticks([-2, -1, 0, 1, 2])
#    ax.set_yticks([-2, -1, 0, 1, 2])


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


    ax.plot(data_x[num*num_grid:(num*num_grid+num_grid)], data_tmp[num*num_grid:(num*num_grid+num_grid)],  "-", color="r", label="Grid")
    ax.plot(data_xM[num*num_grid:(num*num_grid+num_grid)], data_tmpM[num*num_grid:(num*num_grid+num_grid)],  "-", color="b", label="Montecalro")
    ax.legend(loc="upper left")


ani = FuncAnimation(fig, plot, frames=num_snapshots, interval=150)

#print(frames)
ani.save("conduction1D.gif", writer="pillow")
#ani.save('conduction1D.mp4', writer="ffmpeg")

plt.show()