# See http://matplotlib.org/1.4.3/examples/animation/moviewriter.html


from numpy import *
import matplotlib
#matplotlib.use("Agg")
import matplotlib.pyplot as plt
import sys
import matplotlib.animation as manimation
import re

import numpy as np
from matplotlib.animation import FuncAnimation

try:
	sol_filename = sys.argv[1]
except:
	print("Usage: python sol_movie.py <name of the file containing data for u>")
	sys.exit()


rgx = re.search("(.*\/)?(.*\.txt)", sol_filename)
path = "" if rgx.group(1) is None else rgx.group(1)
name = rgx.group(2)

u=loadtxt(sol_filename)
t = loadtxt(path + "time" + name[1:])
x = linspace(0,5, u.shape[1])
Nf = u.shape[0]
j=1


"""
FFMpegWriter = manimation.writers['ffmpeg']
metadata = dict(title=sol_filename, artist='Burgers',
        comment='Look at the solution go!')
writer = FFMpegWriter(fps=15, metadata=metadata)

fig = plt.figure()
l, = plt.plot([], [], 'k-o')


plt.xlim([0,5])
plt.ylim([amin(u),amax(u)])
plt.xlabel('$x$')
plt.ylabel('$u(x)$')


with writer.saving(fig, sol_filename.replace('.txt','.mp4'), 100):
    for i in range(0, Nf, int(Nf/100)):
        plt.title("T = " + str(t[i]))
        l.set_data(x, u[i,:])
        writer.grab_frame()

"""
fig = plt.figure()
ax = fig.add_subplot(111,xlim=(-1,6),ylim=(-1,3))
ax.grid()

line, = ax.plot([],[],'-',lw=2)
time_template = 'time = %.3fs'
time_text = ax.text(0.05,0.9,'',transform=ax.transAxes)

def animate(i):
    line.set_data(x,u[i])
    time_text.set_text(time_template % (t[i]))
    return line, time_text

ani = FuncAnimation(
        fig, animate, len(u),interval=100)
plt.show()

