# -*- coding: utf-8 -*-
"""
Created on Tue May 22 16:48:42 2018

@author: Armando
"""

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from matplotlib.animation import FFMpegWriter


plt.ioff()

metadata = dict(title='Movie Test', artist='Me',
                comment='Movie support!')
writer = FFMpegWriter(fps=2, metadata=metadata)
                     
def f(x, y):
    return np.sin(x) + np.cos(y)

x = np.linspace(0, 2 * np.pi, 120)
y = np.linspace(0, 2 * np.pi, 100).reshape(-1, 1)
# ims is a list of lists, each row is a list of artists to draw in the
# current frame; here we are just animating one artist, the image, in
# each frame
frames = []


for i in range(60):
    x += np.pi / 15.
    y += np.pi / 20.
    curVals = f(x, y)
    frames.append(curVals)

fig = plt.figure()
ax = fig.add_subplot(111)
 
cv0 = frames[0]
im = ax.imshow(cv0, origin='lower',animated =True) # Here make an AxesImage rather than contour
cb = fig.colorbar(im)
tx = ax.set_title('Frame 0')    

def animate(i):
    arr = frames[i]
    vmax     = np.max(arr)
    vmin     = np.min(arr)
    im.set_data(arr)
    im.set_clim(vmin, vmax)
    im.set_cmap(plt.get_cmap('viridis'))
    tx.set_text('Frame {0}'.format(i))
    plt.savefig('C:\\Users\\Armando\\Desktop\\colorbartest\\' +str(i)+ ".png",bbox_inches='tight')
    # In this version you don't have to do anything to the colorbar,
    # it updates itself when the mappable it watches (im) changes

ani = animation.FuncAnimation(fig, animate, frames=len(frames),interval=500,  repeat = False)

writer = FFMpegWriter(fps=2, metadata=dict(artist='Me'), bitrate=1800)
ani.save("C:\\Users\\Armando\\Desktop\\colorbartest\\movie.mp4", writer=writer)

#with writer.saving(fig, "C:\\Users\\Armando\\Desktop\\colorbartest\\movie.mp4", 100):
#    for i in range(100):
#        arr = frames[i]
#        vmax = np.max(arr)
#        vmin = np.min(arr)
#        im.set_data(arr)
#        im.set_clim(vmin, vmax)
#        im.set_cmap(plt.get_cmap('viridis'))
#        tx.set_text('Frame {0}'.format(i))
#        writer.grab_frame()

#plt.show()