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

# データ取得
data_x, data_y, data_vx, data_vy = np.loadtxt("./output_pos.dat", comments='#', unpack=True)

# 粒子数、スナップショット数取得
num_p = 1000
num_tot = len(data_x)
num_snapshots = int(num_tot/num_p)

# 図、枠（軸）の設定
fig = plt.figure()
ax1 = fig.add_subplot(121)
ax2 = fig.add_subplot(122)

def plot(num):
    # clear previous frames
    ax1.cla()
    # Axesのタイトルを設定
    ax1.set_title("2D Maxwell-Boltzmann (XY)", size = 12)

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

    # 軸目盛を設定
    ax1.set_xticks([0, 0.2, 0.4, 0.6, 0.8, 1])
    ax1.set_yticks([0, 0.2, 0.4, 0.6, 0.8, 1])

    # 軸範囲の設定
    ax1.set_xlim(-0.2, 1.2)
    ax1.set_ylim(-0.2, 1.2)

    ax1.scatter(data_x[num*num_p:(num*num_p+num_p)], data_y[num*num_p:(num*num_p+num_p)], s = 1, color = "red")

    # clear previous frames
    ax2.cla()
    # Axesのタイトルを設定
    ax2.set_title("2D Maxwell-Boltzmann (VxVy)", size = 12)

    # 軸ラベルを設定
    ax2.set_xlabel("$v_x$", size = 14, color = "b")
    ax2.set_ylabel("$v_y$", size = 14, color = "b",labelpad = -10)

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

    # 軸範囲の設定
    ax2.set_xlim(-2, 2)
    ax2.set_ylim(-2, 2)

    ax2.scatter(data_vx[num*num_p:(num*num_p+num_p)], data_vy[num*num_p:(num*num_p+num_p)], s = 1, color = "blue")



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

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

plt.show()