#!/usr/bin/env python
# coding: utf-8

# In[1]:


import numpy as np
import matplotlib.pyplot as plt
from salishsea_tools import viz_tools, places, visualisations
import netCDF4 as nc # unless you prefer xarray
import datetime as dt
import os
import glob
import cmocean
from cartopy import crs, feature
import matplotlib
import xarray as xr

get_ipython().run_line_magic('matplotlib', 'inline')


# In[2]:


plt.rcParams.update({'font.size': 12, 'axes.titlesize': 'medium'})


# In[3]:


# load model mesh
with nc.Dataset('/data/eolson/results/MEOPAR/NEMO-forcing-new/grid/mesh_mask201702.nc') as fm:
    print(fm.variables.keys())
    tmask=fm.variables['tmask'][:,:,:,:]
    navlon=fm.variables['nav_lon'][:,:]
    navlat=fm.variables['nav_lat'][:,:]


# In[70]:


bathy = xr.open_dataset('/home/sallen/MEOPAR/grid/bathymetry_201702.nc')

runoff1 = xr.open_dataset('/results/forcing/rivers/R201702DFraCElse_y2019m06d01.nc')
runoff2 = xr.open_dataset('/results/forcing/rivers/R201702DFraCElse_y2019m12d01.nc')

CYrunoff = xr.open_dataset('/results/forcing/rivers/R201702DFraCElse_y2008m06d01.nc')
WYrunoff = xr.open_dataset('/results/forcing/rivers/R201702DFraCElse_y2019m06d01.nc')

mesh = xr.open_dataset('/home/sallen/MEOPAR/grid/mesh_mask201702.nc')
# the mask is 1 where there is water, we want the opposite.  The meshmask has an extra dimension, hence the [0]
tmask = 1 - mesh.tmask[0]


# In[71]:


cmap0=cmocean.cm.deep
cmap0.set_bad('lightgray')

# exclude Fraser off the edge
jmax = 398
jj = range(jmax)
ii = range(898)
jjm, iim = np.meshgrid(jj, ii)
fluxarray = np.array(runoff1.rorunoff[0, :, :jmax])
fig, ax = plt.subplots(1, 1, figsize=(5, 9))
colours = ax.contourf(bathy.Bathymetry[:, :jmax], cmap=cmap0)


ax.set_xlabel('grid X')
ax.set_ylabel('grid Y')
# ax.set_title('Bathymetry of SalishSeaCast')


ax.scatter(jjm[fluxarray>0], iim[fluxarray>0], s=fluxarray[fluxarray>0]*1000, color='royalblue',alpha=0.5)
ax.set_title('River Discharge [kg m-2 s-1]\nJune 1, 2019')




# cb = fig.colorbar(colours, ax=ax)
# cb.set_label('Seafloor Depth (m)');
viz_tools.set_aspect(ax);


# In[72]:


cmap0=cmocean.cm.deep
cmap0.set_bad('lightgray')

# exclude Fraser off the edge
jmax = 398
jj = range(jmax)
ii = range(898)
jjm, iim = np.meshgrid(jj, ii)
fluxarray = np.array(runoff2.rorunoff[0, :, :jmax])
fig, ax = plt.subplots(1, 1, figsize=(5, 9))
colours = ax.contourf(bathy.Bathymetry[:, :jmax], cmap=cmap0)

ax.set_xlabel('grid X')
ax.set_ylabel('grid Y')
# ax.set_title('Bathymetry of SalishSeaCast')


ax.scatter(jjm[fluxarray>0], iim[fluxarray>0], s=fluxarray[fluxarray>0]*1000, color='royalblue',alpha=0.5)
ax.set_title('River Discharge [kg m-2 s-1]\nDecember 1, 2019')




# cb = fig.colorbar(colours, ax=ax)
# cb.set_label('Seafloor Depth (m)');
viz_tools.set_aspect(ax);


# In[69]:


cmap0=cmocean.cm.deep
cmap0.set_bad('lightgray')

# exclude Fraser off the edge
jmax = 398
jj = range(jmax)
ii = range(898)
jjm, iim = np.meshgrid(jj, ii)
fluxarray = np.array(runoff1.rorunoff[0, :, :jmax])
fluxarray2 = np.array(runoff2.rorunoff[0, :, :jmax])

fig, ax = plt.subplots(1, 2, figsize=(8, 8))
colours = ax[0].contourf(bathy.Bathymetry[:, :jmax], cmap=cmap0)


ax[0].set_xlabel('grid X')
ax[0].set_ylabel('grid Y')
# ax.set_title('Bathymetry of SalishSeaCast')



ax[0].scatter(jjm[fluxarray>0], iim[fluxarray>0], s=fluxarray[fluxarray>0]*1000, color='royalblue',alpha=0.5)
ax[0].set_title('June 1, 2019\nRiver Discharge (kg m$^{-2}$ s$^{-1}$)')




# cb = fig.colorbar(colours, ax=ax)
# cb.set_label('Seafloor Depth (m)');
viz_tools.set_aspect(ax[0]);



colours = ax[1].contourf(bathy.Bathymetry[:, :jmax], cmap=cmap0)

ax[1].set_xlabel('grid X')
ax[1].set_ylabel('grid Y')
# ax.set_title('Bathymetry of SalishSeaCast')


ax[1].scatter(jjm[fluxarray2>0], iim[fluxarray2>0], s=fluxarray2[fluxarray2>0]*1000, color='royalblue',alpha=0.5)
ax[1].set_title('December 1, 2019\nRiver Discharge (kg m$^{-2}$ s$^{-1}$)')




# cb = fig.colorbar(colours, ax=ax)
# cb.set_label('Seafloor Depth (m)');
viz_tools.set_aspect(ax[1]);

#plt.savefig('SuppFigureS1_Rivers.png', bbox_inches='tight',dpi=1000,transparent=False)


# In[73]:


cmap0=cmocean.cm.deep
cmap0.set_bad('lightgray')

# exclude Fraser off the edge
jmax = 398
jj = range(jmax)
ii = range(898)
jjm, iim = np.meshgrid(jj, ii)
fluxarray = np.array(CYrunoff.rorunoff[0, :, :jmax])
fluxarray2 = np.array(WYrunoff.rorunoff[0, :, :jmax])

fig, ax = plt.subplots(1, 2, figsize=(8, 8))
colours = ax[0].contourf(bathy.Bathymetry[:, :jmax], cmap=cmap0)


ax[0].set_xlabel('grid X')
ax[0].set_ylabel('grid Y')
# ax.set_title('Bathymetry of SalishSeaCast')



ax[0].scatter(jjm[fluxarray>0], iim[fluxarray>0], s=fluxarray[fluxarray>0]*1000, color='royalblue',alpha=0.5)
ax[0].set_title('Typical Cold Year\nRiver Discharge (kg m$^{-2}$ s$^{-1}$)')




# cb = fig.colorbar(colours, ax=ax)
# cb.set_label('Seafloor Depth (m)');
viz_tools.set_aspect(ax[0]);



colours = ax[1].contourf(bathy.Bathymetry[:, :jmax], cmap=cmap0)

ax[1].set_xlabel('grid X')
ax[1].set_ylabel('grid Y')
# ax.set_title('Bathymetry of SalishSeaCast')


ax[1].scatter(jjm[fluxarray2>0], iim[fluxarray2>0], s=fluxarray2[fluxarray2>0]*1000, color='royalblue',alpha=0.5)
ax[1].set_title('Typical Warm Year\nRiver Discharge (kg m$^{-2}$ s$^{-1}$)')




# cb = fig.colorbar(colours, ax=ax)
# cb.set_label('Seafloor Depth (m)');
viz_tools.set_aspect(ax[1]);

#plt.savefig('SuppFigureS1_Rivers.png', bbox_inches='tight',dpi=1000,transparent=False)


# In[ ]: