"""
author: Jake Vanderplas, email: vanderplas@astro.washington.edu
website: http://jakevdp.github.com, license: BSD
Please feel free to use and modify this, but keep the above information. Thanks!
"""

import numpy as np
from matplotlib import pyplot as plt
from matplotlib import animation
from matplotlib.animation import PillowWriter
from matplotlib.ticker import (MultipleLocator, AutoMinorLocator)

# First set up the figure, the axis, and the plot element we want to animate
fig = plt.figure()
dX = 10   #Set X-axis range
dY = 5.0  #Set Y-axis half range
dW = 2.0  #Set line width
ax = plt.axes(xlim=(0, dX), ylim=(-dY*1.2, dY*1.2))

line, = ax.plot([], [], lw=dW)
ax.xaxis.set_major_locator(MultipleLocator(dX/5))
ax.xaxis.set_major_formatter('{x:.1f}')
#ax.xaxis.set_minor_locator(MultipleLocator(dX/100))
ax.xaxis.set_minor_locator(AutoMinorLocator())
ax.tick_params(which='both', width=2)
ax.tick_params(which='major', length=2)
ax.tick_params(which='minor', length=1, color='r')
ax.grid()
ax.yaxis.set_major_locator(MultipleLocator(dY/5))
ax.yaxis.set_major_formatter('{x:.1f}')
ax.yaxis.set_minor_locator(AutoMinorLocator())

ax.grid(axis="x", color="green", alpha=0.2, linewidth=0.5, linestyle=":")
ax.grid(axis="y", color="black", alpha=0.2, linewidth=0.5, linestyle=":")
ax.tick_params(axis="x", labelsize=10, labelrotation=45, labelcolor="red")
ax.tick_params(axis="y", labelsize=10, labelrotation=45, labelcolor="orange")

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

# animation function.  This is called sequentially
def animate(i):
    x = np.linspace(0, dX, 1000)
    y = dY * np.sin(np.pi * (x - 0.01 * i))
    line.set_data(x, y)
    return line,

#blit=True means only re-draw the parts that have changed. Delay between frames
#in milliseconds. Defaults to 200. animate:function to call at each frame
#Ref: matplotlib.org/3.1.1/api/_as_gen/matplotlib.animation.FuncAnimation.html
anim = animation.FuncAnimation(fig, animate, init_func=init,
    frames=200, interval=10, blit=True)

#Save the animation as an mp4.  This requires ffmpeg or mencoder
#More info: matplotlib.sourceforge.net/api/animation_api.html
#anim.save('movingWave.mp4', fps=30, extra_args=['-vcodec', 'libx264'])

#Save animation as GIF using Pillow module
anim.save("movingWave.gif", writer=PillowWriter(fps=24))

plt.show()