import numpy as np
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,)

tracers1 = nc.Dataset('/data/nsoontie/MEOPAR/SalishSea/results/storm-surges/tide_fix/dec2006/all_forcing/1hour/SalishSea_1h_20061214_20061215_grid_T.nc')

tracers2 = nc.Dataset('/data/nsoontie/MEOPAR/SalishSea/results/storm-surges/tide_fix/dec2006/tidesonly/1hour/SalishSea_1h_20061214_20061215_grid_T.nc')

timesteps = tracers1.variables['time_counter']

ssh1 = tracers1.variables['sossheig']
ssh2 = tracers2.variables['sossheig']

ssh3 = np.zeros_like(ssh1)
ssh3 = np.subtract(ssh1, ssh2)

#Load the data
path='/data/nsoontie/MEOPAR/SalishSea/results/storm-surges/tide_fix/dec2006/'

runs = {'all_forcing/1hour', 'tidesonly/1hour'}

fUs={}; fVs={}; fTs={};

for key in runs:
    fUs[key] = nc.Dataset(path + key +'/SalishSea_1h_20061214_20061215_grid_U.nc','r');
    fVs[key] = nc.Dataset(path + key +'/SalishSea_1h_20061214_20061215_grid_V.nc','r');
    fTs[key] = nc.Dataset(path + key +'/SalishSea_1h_20061214_20061215_grid_T.nc','r');
    
runs = {'all_forcing','tidesonly'}

for key in runs:
    fUs[key] = nc.Dataset(path + key +'/SalishSea_4h_20061211_20061217_grid_U.nc','r');
    fVs[key] = nc.Dataset(path + key +'/SalishSea_4h_20061211_20061217_grid_V.nc','r');
    fTs[key] = nc.Dataset(path + key +'/SalishSea_4h_20061211_20061217_grid_T.nc','r');

t=41
Us={}; Vs={}; Es ={}; Ss={}; Ts={};
depthlevel=0

for key in {'all_forcing', 'tidesonly'}:
    [Us[key], Vs[key], Es[key], Ss[key], Ts[key]] = stormtools.get_variables(fUs[key],fVs[key],fTs[key],t,depthlevel)
background = Es['all_forcing']-Es['tidesonly'];

ssh3_back = np.zeros_like(ssh3)
ssh3_back = np.subtract(ssh3, background)

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

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

#Making an initial image i.e. our first ssh reading
def init():
    ssh = np.ma.masked_values(ssh3_back[0], 0)
    mesh = ax.pcolormesh(ssh,cmap=cmap)
    mesh.set_clim(cmin,cmax)
    cbar = fig.colorbar(mesh)
    cbar.set_label('{label} [{units}]'.format(label=ssh1.long_name.title(), units=ssh1.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(ssh3_back[t], 0)
    mesh = ax.pcolormesh(ssh,cmap=cmap)
    mesh.set_clim(cmin,cmax)
    ax.set_xlabel('x Index')
    ax.set_ylabel('y Index')
    ax.set_title('t = {:.1f}h'.format((timesteps[t] / 3600)-72))

#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('StormSurge_Animation.mp4',writer=mywriter,fps=1)

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