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

# データ取得
data_E, data_num = np.loadtxt("./output_E.dat", comments='#', unpack=True)

# グリッド数、スナップショット数取得
num_grid = 99
num_tot = len(data_E)
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("Maxwell-Boltzmann dist.", size = 10)

    # 軸ラベルを設定
    ax.set_xlabel("$E$", size = 14, color = "k")
    ax.set_ylabel("$dN/dE$", 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_E[num*num_grid:(num*num_grid+num_grid)], data_num[num*num_grid:(num*num_grid+num_grid)],  "-", color="r", label="Montecalro")
    ax.legend(loc="upper right")


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

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

plt.show()