import netCDF4 as NC
import numpy as np

import matplotlib.pyplot as plt
import matplotlib.cm as cm

from salishsea_tools import (nc_tools,viz_tools,tidetools)

%matplotlib inline

tracers1 = NC.Dataset('/ocean/nsoontie/MEOPAR/SalishSea/results/storm-surges/final/dec2006/all_forcing/30min/SalishSea_30m_20061214_20061215_grid_T.nc')
nc_tools.show_variables(tracers1)
ssh1=tracers1.variables['sossheig']
sal1=tracers1.variables['vosaline']
temp1=tracers1.variables['votemper']
zlevel1 = tracers1.variables['deptht']
timesteps1 = tracers1.variables['time_counter']

tracers2 = NC.Dataset('/data/dlatorne/MEOPAR/SalishSea/nowcast/24-26sep14/SalishSea_1h_20140924_20140926_grid_T.nc')
nc_tools.show_variables(tracers2)
ssh2=tracers2.variables['sossheig']
sal2=tracers2.variables['vosaline']
temp2=tracers2.variables['votemper']
zlevel2 = tracers2.variables['deptht']
timesteps2 = tracers2.variables['time_counter']

temps = [temp1,temp2]
titles = ['December 2006','September 2014']
fig, axs = plt.subplots(1, 2, figsize=(10, 8))

for ax,temp,title in zip(axs,temps,titles):
    t, z = 0, 0
    temp_tz = np.ma.masked_values(temp[t, z], 0)
    viz_tools.set_aspect(ax)
    cmap = plt.get_cmap('hsv')
    cmap.set_bad('burlywood')
    mesh = ax.pcolormesh(temp_tz, cmap=cmap,vmin=0,vmax=18)
    cbar = fig.colorbar(mesh,ax=ax)
    plt.axis((0, temp_tz.shape[1], 0, temp_tz.shape[0]))

    ax.grid()

    ax.set_xlabel('x Index')
    ax.set_ylabel('y Index')
    cbar.set_label('{label} [{units}]'.format(label=temp.long_name.title(), units=temp.units))
    ax.set_title(title)


temps = [temp1,temp2]
titles = ['December 2006','September 2014']
fig, axs = plt.subplots(1, 2, figsize=(20, 12))

for ax,temp,title in zip(axs,temps,titles):
    t, z = 0, 0
    temp_tz = np.ma.masked_values(temp[t, z], 0)
    viz_tools.set_aspect(ax)
    cmap = plt.get_cmap('gist_ncar_r')
    cmap.set_bad('burlywood')
    mesh = ax.pcolormesh(temp_tz, cmap=cmap)
    cbar = fig.colorbar(mesh,ax=ax)
    ax.set_xlim(250,temp_tz.shape[1])
    ax.set_ylim(350, 550)

    ax.set_xlabel('x Index')
    ax.set_ylabel('y Index')
    cbar.set_label('{label} [{units}]'.format(label=temp.long_name.title(), units=temp.units))
    ax.set_title(title)
    
    ax.annotate('Fraser',(330,420),fontsize=14,color='white')
    ax.annotate('Squamish',(360,540),fontsize=14,color='white')
    
    plt.suptitle('Note: Colorbar range varies for subplots',fontsize=15)
    
#    
fig, axs = plt.subplots(1, 2, figsize=(20, 12))

for ax,temp,title in zip(axs,temps,titles):
    t, z = 0, 0
    temp_tz = np.ma.masked_values(temp[t, z], 0)
    viz_tools.set_aspect(ax)
    cmap = plt.get_cmap('gist_ncar_r')
    cmap.set_bad('burlywood')
    mesh = ax.pcolormesh(temp_tz, cmap=cmap)
    cbar = fig.colorbar(mesh,ax=ax)
    ax.set_xlim(50,250)
    ax.set_ylim(700, 900)

    ax.set_xlabel('x Index')
    ax.set_ylabel('y Index')
    cbar.set_label('{label} [{units}]'.format(label=temp.long_name.title(), units=temp.units))
    ax.set_title(title)
    
    ax.annotate('Campbell',(80,750),fontsize=14,color='white')
    
    plt.suptitle('Note: Colorbar range same as above',fontsize=15)