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

# In[1]:


import matplotlib.pyplot as plt
import netCDF4 as nc
import numpy as np
import xarray as xr

from salishsea_tools import visualisations as vis
from salishsea_tools import viz_tools



# In[4]:


dataset = xr.open_dataset('/ocean/jtao/MOAD/extractions/SalishSeaCast_wastewater_month_outfall_20180101_20190930.nc')
dataset


# In[5]:


meshmask = xr.open_dataset('/home/sallen/MEOPAR/grid/mesh_mask202108.nc')
tmask = meshmask.tmask[0]
tmask = tmask.rename({'z' : 'depth', 'y': 'gridY', 'x': 'gridX'})
tmask


# In[6]:


mesh = nc.Dataset('/home/sallen/MEOPAR/grid/mesh_mask202108.nc')
bathy = nc.Dataset('/home/sallen/MEOPAR/grid/bathymetry_202108.nc')
mesh_mask_depth_var='gdept_0'
thalweg_file = '/home/sallen/MEOPAR/Tools/bathymetry/thalweg_working.txt'
depth = mesh.variables[mesh_mask_depth_var][:]
thalweg_pts = np.loadtxt(thalweg_file, delimiter=' ', dtype=int)


# In[7]:


fig, axs = plt.subplots(6, 4, figsize=(20, 6*9))
ip = 0
jp = 0
for month in range(1, 12+9+1):
    dataset.outfall[month-1, 0].where(tmask[0]).plot(ax=axs[ip, jp], vmax=2e-4);
    jp = jp + 1;
    if jp == 4:
        ip = ip + 1
        jp = 0
fig.tight_layout();


# In[9]:


fig, axs = plt.subplots(21, 1, figsize=(20, 72))
jp = 0
for month in range(1, 12+9+1):
    dep_thal, distance, var_thal = vis.load_thalweg(
                depth[0, ...], np.array(dataset.outfall[month-1]), bathy['nav_lon'][:], bathy['nav_lat'][:],
                thalweg_pts)
    var_plot = vis._fill_in_bathy(var_thal, mesh, thalweg_pts)
    colours = axs[jp].contourf(distance, dep_thal, var_plot, np.arange(0, 4e-4, 1e-5), 
                                     cmap='plasma', extend='both', zorder=-1)
    axs[jp].contour(distance, dep_thal, var_plot, [2e-5, 5e-5, 1e-4, 2e-4], 
                                     colors='green', zorder=0)
    vis._add_bathy_patch(
            distance, bathy['Bathymetry'][:], thalweg_pts, axs[jp], color='burlywood')
    cbar = fig.colorbar(colours, ax=axs[jp])
    axs[jp].invert_yaxis()
    axs[jp].set_ylim(200, 0)
    axs[jp].set_ylabel('Depth [m]');
    axs[jp].set_title(f'{month} month of 2018')
    jp = jp + 1;
fig.tight_layout();


# In[29]:


245/12


# In[31]:


20.42/30 * 9.5


# In[35]:


6.66*4*2 + 14


# In[36]:


67.28/245


# In[ ]: