import g3read as g3
import matplotlib as mpl
mpl.use('Agg')
import matplotlib.pyplot as plt
from matplotlib import animation
plt.rcParams['animation.ffmpeg_path'] = '/usr/bin/ffmpeg'


"""
Tutorial here: https://jakevdp.github.io/blog/2012/08/18/matplotlib-animation-tutorial/
"""

# Set up formatting for the movie files
Writer = animation.writers['ffmpeg']
writer = Writer(fps=15, metadata=dict(artist='Me'), bitrate=1800)

# First set up the figure, the axis, and the plot element we want to animate
fig = plt.figure()
ax = plt.axes(xlim=(0, 1), ylim=(0, 1))
ax.set_aspect('equal')
point, = ax.plot([], [], 'o')

# initialization function: plot the background of each frame
def init():
    point.set_data([], [])
    return point,

# animation function.  This is called sequentially
def animate(i):
    filename = f'./snap_{i:03}'

    f = g3.GadgetFile(filename)
    pos = f.read_new("POS ", 0)
    point.set_data(pos[:,0], pos[:,1])
    return point,

# call the animator.  blit=True means only re-draw the parts that have changed.
anim = animation.FuncAnimation(fig, animate, init_func=init,
                               frames=11, blit=True)

# save the animation as an mp4.  This requires ffmpeg or mencoder to be
# installed.  The extra_args ensure that the x264 codec is used, so that
# the video can be embedded in html5.  You may need to adjust this for
# your system: for more information, see
# http://matplotlib.sourceforge.net/api/animation_api.html
anim.save('basic_animation.mp4',  writer=writer)