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_z = np.loadtxt("./output.dat", comments='!', unpack=True)

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

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


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

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

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

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

    ax.scatter(data_x[num*num_p:(num*num_p+num_p)], data_y[num*num_p:(num*num_p+num_p)], data_z[num*num_p:(num*num_p+num_p)], s = 1, color = "blue")



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

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

plt.show()