from itertools import count
import pandas as pd
import matplotlib
matplotlib.use("Qt4Agg")
import matplotlib.pyplot as plt
#from scipy.integrate import solve_ivp, odeint
from matplotlib.animation import FuncAnimation
import numpy as np
from math import pi
# Set up formatting for the movie files
Writer = matplotlib.animation.writers['ffmpeg']
writer = Writer(fps=15, metadata=dict(artist='Me'), bitrate=1800)


def xc(r, n, t):
    return r * np.cos(n*t)

def yc(r, n, t):
    return r * np.sin(n*t)

    
r = 1.
T = 360
n_steps = 360
t = np.linspace(0. ,T , n_steps)/T*2*pi

nx = 1.
ny = 1.

#set up the plot
fig, ax = plt.subplots()
d = 1.2*r
ax.set_xlim(-d, d)
ax.set_ylim(-d, d)
line, = ax.plot(xc(r, nx, t)[0], yc(r, ny, t)[0])
ax.set_title(r'$(x, y) = (\cos({} t), \sin({} t)$'.format(nx, ny))
#line, = ax.plot(xc(r, nx, t), yc(r, ny, t))
plt.gca().set_aspect('equal', adjustable='box')
ax.set_xlabel('x',fontsize=14)
ax.set_ylabel('y',fontsize=14)
#plt.show()

x = xc(r, nx, t)
y = yc(r, ny, t)

xvals = []
yvals = []
index = count()

line, = ax.plot([], [], 'k-')
pos, = ax.plot([], [], 'ko')

def animate(i):
    #print(x[i],y[i])
    xvals.append(x[i])
    yvals.append(y[i])
    pos.set_xdata(x[i])
    pos.set_ydata(y[i])
    line.set_xdata(xvals)
    line.set_ydata(yvals)
    return line

ani = FuncAnimation(fig, func=animate, frames=np.arange(0,n_steps-1, 1), interval = 100)

plt.tight_layout()
#plt.legend(loc='best')
plt.gca().set_aspect('equal', adjustable='box')

try:
    writer = matplotlib.animation.writers['avconv']
except KeyError:
    writer = matplotlib.animation.writers['ffmpeg']
writer = writer(fps=12)
filename = 'wave-dispersion-animation.mp4'.format(nx, ny)

ani.save(filename, writer=writer, dpi=300)
plt.close(fig)