import numpy as np
import matplotlib.cm as cm
import matplotlib.mlab as mlab
from matplotlib import pyplot as plt
from matplotlib import animation
from salishsea_tools import stormtools

import netCDF4 as nc
from salishsea_tools import (
    nc_tools,
    viz_tools,)

%matplotlib inline

tracers = nc.Dataset('/ocean/imachuca/MEOPAR/data/results/plume/jun2003/SalishSea_1h_20030531_20030602_grid_T.nc')

nc_tools.show_variables(tracers)

timesteps = tracers.variables['time_counter']

bathy = nc.Dataset('/ocean/imachuca/MEOPAR/NEMO-forcing/grid/bathy_meter_SalishSea2.nc')

salvar = tracers.variables['vosaline']
salvar.shape

sal0 = salvar[:,0,:,:]

cmin = sal0.min()
cmax = sal0.max()

#Setting up a blank figure
fig, ax = plt.subplots(1, 1, figsize=(10, 8))
viz_tools.set_aspect(ax)
cmap = plt.get_cmap('Blues')
viz_tools.plot_land_mask(ax, bathy, color='burlywood')

#Making an initial image i.e. our first ssh reading
def init():
    sal = np.ma.masked_values(sal0[0], 0)
    mesh = ax.pcolormesh(sal,cmap=cmap)
    mesh.set_clim(cmin,cmax)
    cbar = fig.colorbar(mesh)
    cbar.set_label('{label} [{units}]'.format(label=salvar.long_name.title(), units=salvar.units))
    ax.set_xlabel('x Index')
    ax.set_ylabel('y Index')
    ax.set_title('t = 0.0 h')

    
#The full range of images that will make up the animation
def animate(t):
    ssh = np.ma.masked_values(sal0[t], 0)
    mesh = ax.pcolormesh(ssh,cmap=cmap)
    mesh.set_clim(cmin,cmax)
    ax.set_xlabel('x Index')
    ax.set_ylabel('y Index')
    timestamp = nc_tools.timestamp(tracers,t)
    ax.set_title(timestamp)

#The animation function    
anim = animation.FuncAnimation(fig, animate, init_func=init, frames=48, interval=300, blit=False, repeat=False)

#A line that makes it all work
mywriter = animation.FFMpegWriter()

#Save in current folder
anim.save('SalinityJun_Animation.mp4',writer=mywriter,fps=1)

#Show as a pop-up window
#plt.show()


#Setting up a blank figure
fig, ax = plt.subplots(1, 1, figsize=(10, 8))
viz_tools.set_aspect(ax)
cmap = plt.get_cmap('Blues')
viz_tools.plot_land_mask(ax, bathy, color='burlywood')

#Making an initial image i.e. our first ssh reading
def init():
    sal = np.ma.masked_values(sal0[0], 0)
    mesh = ax.pcolormesh(sal,cmap=cmap)
    mesh.set_clim(cmin,cmax)
    cbar = fig.colorbar(mesh)
    cbar.set_label('{label} [{units}]'.format(label=salvar.long_name.title(), units=salvar.units))
    ax.set_xlabel('x Index')
    ax.set_ylabel('y Index')
    ax.set_title('t = 0.0 h')

    
#The full range of images that will make up the animation
def animate(t):
    plt.cla()
    ssh = np.ma.masked_values(sal0[t], 0)
    mesh = ax.pcolormesh(ssh,cmap=cmap)
    
    cont = ax.contour(ssh,5, colors='k')
    class nf(float):
        def __repr__(self):
         str = '%.1f' % (self.__float__(),)
         if str[-1]=='0':
             return '%.0f' % self.__float__()
         else:
             return '%.1f' % self.__float__()
    cont.levels = [nf(val) for val in cont.levels ]
    fmt=str
    ax.clabel(cont, cont.levels, fmt = fmt, inline=True, fontsize=9)
    
    mesh.set_clim(cmin,cmax)
    ax.set_xlabel('x Index')
    ax.set_ylabel('y Index')
    timestamp = nc_tools.timestamp(tracers,t)
    ax.set_title(timestamp)
    

#The animation function    
anim = animation.FuncAnimation(fig, animate, init_func=init, frames=24, interval=300, blit=False, repeat=False)

#A line that makes it all work
mywriter = animation.FFMpegWriter()

#Save in current folder
anim.save('SalinityJunC_Animation.mp4',writer=mywriter,fps=1)

#Show as a pop-up window
#plt.show()



#Setting up a blank figure
fig, ax = plt.subplots(1, 1, figsize=(10, 8))
viz_tools.set_aspect(ax)
cmap = plt.get_cmap('ocean_r')
viz_tools.plot_land_mask(ax, bathy, color='burlywood')

#Making an initial image i.e. our first ssh reading
def init():
    sal = np.ma.masked_values(sal0[0], 0)
    mesh = ax.pcolormesh(sal,cmap=cmap)
    mesh.set_clim(cmin,cmax)
    cbar = fig.colorbar(mesh)
    cbar.set_label('{label} [{units}]'.format(label=salvar.long_name.title(), units=salvar.units))
    ax.set_xlabel('x Index')
    ax.set_ylabel('y Index')
    ax.set_title('t = 0.0 h')

    
#The full range of images that will make up the animation
def animate(t):
    ssh = np.ma.masked_values(sal0[t], 0)
    mesh = ax.pcolormesh(ssh,cmap=cmap)
    mesh.set_clim(cmin,cmax)
    ax.set_xlabel('x Index')
    ax.set_ylabel('y Index')
    timestamp = nc_tools.timestamp(tracers,t)
    ax.set_title(timestamp)

#The animation function    
anim = animation.FuncAnimation(fig, animate, init_func=init, frames=48, interval=300, blit=False, repeat=False)

#A line that makes it all work
mywriter = animation.FFMpegWriter()

#Save in current folder
anim.save('SalinityJun2_Animation.mp4',writer=mywriter,fps=1)

#Show as a pop-up window
#plt.show()