from __future__ import division, print_function

%matplotlib inline

import matplotlib.pyplot as plt

import numpy as np

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')

nc_tools.show_variables(tracers1)
nc_tools.show_variables(tracers2)

timesteps = tracers1.variables['time_counter']
#timesteps will be used for plot

#Defining sea surface height as ssh1 and ssh2 for "all forcing" and "tides only" respectively
ssh1 = tracers1.variables['sossheig']
ssh2 = tracers2.variables['sossheig']

def ssh_fxn(timeindex):
    
    ssh1_t = np.ma.masked_values(ssh1[timeindex], 0)

    fig, ax = plt.subplots(1, 1, figsize=(10, 8))
    viz_tools.set_aspect(ax)
    cmap = plt.get_cmap('CMRmap')
    cmap.set_bad('burlywood')
    mesh = ax.pcolormesh(ssh1_t, cmap=cmap)
    cbar = fig.colorbar(mesh)

    ax.set_xlabel('x Index')
    ax.set_ylabel('y Index')
    ax.set_title('t = {:.1f}h'.format((timesteps[timeindex] / 3600)-72))
    cbar.set_label('{label} [{units}]'.format(label=ssh1.long_name.title(), units=ssh1.units))

timeindex=1
ssh_fxn(timeindex)

timesteps_comp = (0, 1, 2)
ssh_comp = (ssh1,ssh2)
fig, axs = plt.subplots(1, 3, figsize=(16, 8), sharey=True)


for ax, ssh in zip(axs, ssh_comp): 
    fig, axs = plt.subplots(1, 3, figsize=(16, 8), sharey=True)
    cmap = plt.get_cmap('jet')
    cmap.set_bad('burlywood')
    for ax, t in zip(axs, timesteps_comp):
        ssh_t = np.ma.masked_values(ssh[t], 0)
        viz_tools.set_aspect(ax)
        mesh = ax.pcolormesh(ssh_t, cmap=cmap)
        mesh.set_clim(vmin=-0.75,vmax=0.3)
        cbar = fig.colorbar(mesh, ax=ax)
        ax.set_xlabel('x Index')
    axs[0].set_ylabel('y Index')
        

        
fig, axs = plt.subplots(1, 3, figsize=(16, 8), sharey=True)
cmap = plt.get_cmap('jet')
cmap.set_bad('burlywood')
for ax, t in zip(axs, timesteps_comp):
    ssh3_t = np.ma.masked_values((ssh1[t]-ssh2[t]), 0)
    viz_tools.set_aspect(ax)
    mesh = ax.pcolormesh(ssh3_t, cmap=cmap)
    mesh.set_clim(vmin=-0.75,vmax=0.3)
    cbar = fig.colorbar(mesh, ax=ax)
    ax.set_xlabel('x Index')

axs[0].set_ylabel('y Index')