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

# Compare stratification and mixing under different vertical resolution.
# 
# Susan's Triple tanh vs original

# In[4]:


import netCDF4 as nc
from salishsea_tools import nc_tools
import matplotlib.pyplot as plt
import numpy as np
import os

from nowcast import analyze
from salishsea_tools import tidetools

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


# In[7]:


results = '/data/nsoontie/MEOPAR/SalishSea/results/2Ddomain/3.6/mixing_paper/'

run1 = 'trpl_tanh_noHoll'
run2 = 'base_aug_noHoll'

path={}
path[run1]= os.path.join(results,run1)
path[run2] = os.path.join(results,run2)

print(path[run1])
print(path[run2])

grid = nc.Dataset('/data/nsoontie/MEOPAR/2Ddomain/grid/bathy2D_36.nc')
bathy = grid.variables['Bathymetry'][:]


# In[8]:


Us = {}; Ws = {}; Ts = {}; Ss={}; Diffs={}; Viscs={}; sshs={}; depths={}; xx={}; zz={}; depthsw={}; zzw={}

runs = [run1, run2]
y=5
for run in runs:
    f = nc.Dataset(os.path.join(path[run],'SalishSea_1d_20030819_20030927_grid_U.nc'))
    Us[run] = f.variables['vozocrtx'][:,:,y,:]
    f = nc.Dataset(os.path.join(path[run],'SalishSea_1d_20030819_20030927_grid_W.nc'))
    Ws[run] = f.variables['vovecrtz'][:,:,y,:]
    Viscs[run] = f.variables['vert_eddy_visc'][:,:,y,:]
    Diffs[run] = f.variables['vert_eddy_diff'][:,:,y,:]
    depthsw[run] = f.variables['depthw'][:]
    f = nc.Dataset(os.path.join(path[run],'SalishSea_1d_20030819_20030927_grid_T.nc'))
    Ts[run] = f.variables['votemper'][:,:,y,:]
    Ss[run] = f.variables['vosaline'][:,:,y,:]
    sshs[run] = f.variables['sossheig'][:,y,:]
    depths[run] = f.variables['deptht'][:]
    x = f.variables['nav_lon'][y,:]
    times=np.arange(sshs[run].shape[0])
    xx[run],zz[run] = np.meshgrid(x,-depths[run][:])
    _, zzw[run] = np.meshgrid(x,-depthsw[run][:])  


# Compare vertical grid spacing

# In[9]:


fig, axs = plt.subplots(1,2,figsize=(10,5))

for run in runs:
    spacing=np.zeros(depths[run].shape)
    spacing[0:-1] = depths[run][1:]- depths[run][0:-1] 
    axs[0].plot(spacing, depths[run], 'o-', label=run)
    axs[1].plot(spacing, depths[run], 'o-', label=run)
for ax in axs:
    ax.legend(loc=0)
    ax.set_ylabel('Depth [m]')
    ax.set_xlabel('Grid spacing')
    ax.grid()
axs[0].set_ylim([450,0])
axs[1].set_ylim([50,0])


# Compare fields

# In[10]:


fig, axs = plt.subplots(10,2,figsize=(15,25))
ts = np.arange(0,40,4)
cs = [28, 29, 30, 31,31.2,31.4,31.6,31.8, 32, 32.2, 32.4, 32.6, 32.8, 33, 33.2, 33.4, 33.6, 33.8, 34 ]
for n, run in enumerate(runs):
    for t,ax in zip(ts,axs[:,n].flat):
        Splot = np.ma.masked_values(Ss[run][t,:,:],0)
        mesh=ax.contourf(xx[run],zz[run],Splot,cs,ls='None',extend='both',cmap = 'hsv')
        ax.set_title('Sal: t={} : {}'.format(t, run))
        CS = ax.contour(xx[run],zz[run],Splot, np.arange(31,34.2,.2),colors='k')
        ax.clabel(CS, inline=1, fontsize=10)
        plt.colorbar(mesh,ax=ax)
        ax.set_ylim([-450,0])
        ax.grid()


# * At some times, trpl_tanh has saltier waterover the sills (t=24). At other times, the sill water is fresher (t=16). Most of the time the sill waters to have similar densities.  
# * Density of the water entering the basin is usually fresher in the triple tanh.
# * The basin is shallower in the trpl_tanh case. We have lower resolutio at those depths with the triple tanh. 
# * What are the partial step bathymetry values?
# * Salty water from SJdF (33.6 contour) penetrates further into domain in the old resolution. 
# 
# Could the issue be that the new grid topography isn't smoothing enough near the open boundary? It seems lik there is less salty water right from the boundary.

# In[11]:


fig, axs = plt.subplots(1,2,figsize=(15,3))
t=1
cs = [28, 29, 30, 31,31.2,31.4,31.6,31.8, 32, 32.2, 32.4, 32.6, 32.8, 33, 33.2, 33.4, 33.6, 33.8, 34 ]
for ax, run in zip(axs, runs):
        Splot = np.ma.masked_values(Ss[run][t,:,:],0)
        mesh=ax.contourf(xx[run],zz[run],Splot,cs,ls='None',extend='both',cmap = 'hsv')
        ax.set_title('Sal: t={} : {}'.format(t, run))
        CS = ax.contour(xx[run],zz[run],Splot, np.arange(31,34.2,.2),colors='k')
        ax.clabel(CS, inline=1, fontsize=10)
        plt.colorbar(mesh,ax=ax)
        ax.set_ylim([-450,0])
        ax.grid()
        ax.set_xlim([0,100])
        ax.set_xlabel('x (km)')


# In[12]:


fig, axs = plt.subplots(10,2,figsize=(15,25))
ts = np.arange(0,40,4)

for n, run in enumerate(runs):
    for t,ax in zip(ts,axs[:,n].flat):
        Tplot = np.ma.masked_values(Ts[run][t,:,:],0)
        mesh=ax.contourf(xx[run],zz[run],Tplot,50,vmin=5,vmax=15,ls='None')
        ax.set_title('Temp: t={} : {}'.format(t, run))
        CS = ax.contour(xx[run],zz[run],Tplot, np.arange(8,10.1,.1),colors='k')
        ax.clabel(CS, inline=1, fontsize=10)
        plt.colorbar(mesh, ax=ax)


# In[13]:


print(Ts[run].max())


# There is something wrong with the temperature of the river water.

# In[14]:


fig, axs = plt.subplots(10,2,figsize=(15,25))
ts = np.arange(0,40,4)
for n, run in enumerate(runs):
    for t,ax in zip(ts,axs[:,n].flat):
        Uplot = np.ma.masked_values(Us[run][t,:,:],0)
        mesh=ax.contourf(xx[run],zz[run],Uplot,np.arange(-1.5,1.5,.1),ls='None')
        ax.set_title('U [m/s]: t={} : {}'.format(t, run))
        plt.colorbar(mesh,ax=ax)


# In[15]:


fig, axs = plt.subplots(10,2,figsize=(15,25))
ts = np.arange(0,40,4)

for n, run in enumerate(runs):
    for t,ax in zip(ts,axs[:,n].flat):
        Wplot = np.ma.masked_values(Ws[run][t,:,:],0)
        mesh=ax.contourf(xx[run],zzw[run],Wplot,np.arange(-.01,.011,.001),ls='None',cmap='bwr')
        ax.set_title('W: t={} : {}'.format(t, run))
        plt.colorbar(mesh,ax=ax)


# Mixing parameters

# In[16]:


fig, axs = plt.subplots(10,2,figsize=(15,25))
ts = np.arange(0,40,4)

for n, run in enumerate(runs):
    for t,ax in zip(ts,axs[:,n].flat):
        Dplot = np.ma.masked_values(Diffs[run][t,:,:],0)
        mesh=ax.contourf(xx[run],zzw[run],Dplot,np.arange(0,100),cmap='hot')
        ax.set_title('Diffusivity : t={} : {}'.format(t, run))
        plt.colorbar(mesh,ax=ax)


# In[17]:


fig, axs = plt.subplots(10,2,figsize=(15,25))
ts = np.arange(0,40,4)

for n, run in enumerate(runs):
    for t,ax in zip(ts,axs[:,n].flat):
        Vplot = np.ma.masked_values(Viscs[run][t,:,:],0)
        mesh=ax.contourf(xx[run],zzw[run],Vplot,np.arange(0,100),cmap='hot')
        ax.set_title('Viscosity : t={} : {}'.format(t, run))
        plt.colorbar(mesh,ax=ax)


# ## Area averaged eddy coeffcieints over time
# 
# Area averging in the mixing region (i=300 to 700) 

# In[18]:


for run in runs:
    Diffs[run] = np.ma.masked_values(Diffs[run],0)
    Viscs[run] = np.ma.masked_values(Viscs[run],0)


# In[19]:


def compare_depth_averaged_mixing(ist,ien,d1,d2):
    
    fig1, axs = plt.subplots(1,2,figsize=(15,4))
    

    for run in runs:
        #index of depths
        ind1 = tidetools.find_model_level(d1,depthsw[run])
        ind2 = tidetools.find_model_level(d2,depthsw[run])
        print(run, depthsw[run][ind1], depthsw[run][ind2])
        Diff_davg = analyze.depth_average(Diffs[run][:,ind1:ind2,ist:ien],depthsw[run][ind1:ind2],depth_axis=1)
        print('Max diff:', run, Diff_davg.max())
        axs[0].plot(np.nanmean(Diff_davg,axis=1),label=run)
        axs[0].legend()
        axs[0].set_xlabel('time')
        axs[0].set_ylabel('Diffusivity m^2/s')
    
        Visc_davg = analyze.depth_average(Viscs[run][:,ind1:ind2,ist:ien],depthsw[run][ind1:ind2],depth_axis=1)
        print('Max visc:', run, Visc_davg.max())
        axs[1].plot(np.nanmean(Visc_davg,axis=1),label=run)
        axs[1].legend()
        axs[1].set_xlabel('time')
        axs[1].set_ylabel('Viscosity m^2/s')
    
    fig2, axs = plt.subplots(1,2,figsize=(15,4))
    for run,ax in zip(runs,axs):
        #index of depths
        ind1 = tidetools.find_model_level(d1,depthsw[run])
        ind2 = tidetools.find_model_level(d2,depthsw[run])
        Dplot =Diffs[run][t,:,:]
        mesh=ax.pcolormesh(np.arange(Diffs[run].shape[-1]),depthsw[run],Dplot,vmin=0,vmax=100,cmap='hot')
        ax.set_title('Diffusivity : t={} : {}'.format(t, run))
        cbar = plt.colorbar(mesh,ax=ax)
        cbar.set_label('m^2/s')
        ax.set_ylim([450,0])
        ax.plot([ien,ien],[depthsw[run][ind1],depthsw[run][ind2]], '-b')
        ax.plot([ist,ist],[depthsw[run][ind1],depthsw[run][ind2]], '-b')
        ax.plot([ist,ien],[depthsw[run][ind1],depthsw[run][ind1]], '-b')
        ax.plot([ist,ien],[depthsw[run][ind2],depthsw[run][ind2]], '-b')
        
    return fig1, fig2


# In[20]:


ist=300
ien=700
d1=1
d2=450

fig1, fig2 = compare_depth_averaged_mixing(ist,ien,d1,d2)


# * Higher eddy values in triple tanh case

# ###Surface

# In[21]:


ist=300
ien=700
d1=1
d2=10

fig1, fig2 = compare_depth_averaged_mixing(ist,ien,d1,d2)


# ### Intermediate

# In[22]:


ist=300
ien=700
d1=50
d2=100

fig1, fig2 = compare_depth_averaged_mixing(ist,ien,d1,d2)


# Averaging box is completly different!

# ###Deep

# In[23]:


ist=300
ien=700
d1=100
d2=400
fig1, fig2 = compare_depth_averaged_mixing(ist,ien,d1,d2)


# * The averaging area is slight different again...

# ###Strait of Juan de Fuca

# In[24]:


ist=0
ien=300
d1=1
d2=400
fig1, fig2 = compare_depth_averaged_mixing(ist,ien,d1,d2)


# Smaller eddy coefficients in the Strait of Juan de Fuca. Dissipation in the boundary later seems to behave differently.

# ### Strait of Georgia

# In[25]:


ist=700
ien=1000
d1=1
d2=450
fig1, fig2 = compare_depth_averaged_mixing(ist,ien,d1,d2)


# ### Whole domain

# In[26]:


ist=0
ien=1100
d1=1
d2=450
fig1, fig2 = compare_depth_averaged_mixing(ist,ien,d1,d2)


# Overall, smaller mixing coefficients in the triple tanh case. BUT, in the mixing reigon, they are higher....

# # Stratification Profiles

# In[27]:


i = 900;

fig, axs = plt.subplots(2,2,figsize=(15,8))
ts = np.arange(6,13,1)

for run, ax in zip(runs, axs[0,:].flat):
    for t in zip(ts):
        ax.plot(Ss[run][t,:,i].T,zz[run][:,i],label='t= {}'.format(t[0]),lw=1.5)
        ax.set_title('Salinity profile SoG: {}'.format(run))
        
ts = np.arange(0,27,4)

for run, ax in zip(runs, axs[1,:].flat):
    for t in zip(ts):
        ax.plot(Ss[run][t,:,i].T,zz[run][:,i],label='t= {}'.format(t[0]),lw=1.5)
        ax.set_title('Salinity profile SoG: {}'.format(run))

    
for ax in axs.flat:
    ax.legend(loc=0)
    ax.grid()
    ax.set_xlim([10,35])
    ax.set_ylim([-50,0])
    ax.set_xlabel('Salinity [psu]')
    ax.set_ylabel('Depth [m]')
        
#for t, ax in zip(ts, axs[:,2].flat):
#    ax.pcolormesh(xx,zz,Ss[run1][t,:,:]-Ss[run2][t,:,:],vmin=-1,vm


# There are some differences in the surface layer. 

# In[28]:


i = 900;

fig, axs = plt.subplots(2,2,figsize=(15,8))
ts = np.arange(0,7,1)

for run, ax in zip(runs, axs[0,:].flat):
    for t in zip(ts):
        ax.plot(Ts[run][t,:,i].T,zz[run][:,i],label='t= {}'.format(t[0]),lw=1.5)
        ax.set_title('Temperature  profile SoG: {}'.format(run))

ts = np.arange(0,27,4)

for run, ax in zip(runs, axs[1,:].flat):
    for t in zip(ts):
        ax.plot(Ts[run][t,:,i].T,zz[run][:,i],label='t= {}'.format(t[0]),lw=1.5) 
        ax.set_title('Temperature  profile SoG: {}'.format(run))
    
for ax in axs.flat:    
    ax.legend(loc=0)
    ax.grid()
    ax.set_xlim([5,20])
    ax.set_ylim([-50,0])
    ax.set_xlabel('Temperature [deg C]')
    ax.set_ylabel('Depth [m]')


# Averaging across SoG basin: 400km to 500km

# In[29]:


i = 800; e=1000;

fig, axs = plt.subplots(3,2,figsize=(15,12))
ts = np.arange(0,27,5)
for run in runs:
    Ss[run] = np.ma.masked_values(Ss[run],0)

for t, ax in zip(ts, axs.flat):
    for run in runs:
        ax.plot(np.nanmean(Ss[run][t,:,i:e],axis=1).T,zz[run][:,i],label=run,lw=1.5)
        ax.set_title('Salinty  profile SoG: t = {}'.format(t))
    
for ax in axs.flat:    
    ax.legend(loc=0)
    ax.grid()
    ax.set_xlim([10,34])
    ax.set_ylim([-50,0])
    ax.set_xlabel('Salinity [psu]')
    ax.set_ylabel('Depth [m]')


# ### Comparison of grids

# I should look more closely at the partial steps and grids more carefully
# 
# Bottom cell thickness is set by partial steps in the following way (NEMO book) -  minimum water thickness of a cell partially filled by bathymetry = min(5, 0.2*e3t), where e3t is the reference grid spacing.
# 
# How does the bottom grid spacing compare before and after the application of partial steps?
# 
# In the mesh_mask.nc files - e3t_0 is the three-dimensional grid spacing. So thickness of bottom cell is max(e3t_0*tmask)
# 
# Also, e3t_1d is the 1D reference grid spacing. I think this means before the partial steps is applied.
# 
# Finally, max(gpdet_0*tmask) is the centre location of the bottom cell. Compare among all of my settings.
# 
# First, look at the depth of the bottom cell.

# In[30]:


meshes = {}
y=5
paths= {run1: '/data/nsoontie/MEOPAR/2Ddomain/vertical_resolution/trpl_tanh/mesh_mask.nc',
        run2: '/data/nsoontie/MEOPAR/2Ddomain/grid/mesh_mask.nc',
        'dbl_tanh': '/data/nsoontie/MEOPAR/2Ddomain/vertical_resolution/dbl_tanh/mesh_mask.nc'}
for run in paths:
    f = nc.Dataset(paths[run])
    meshes[run] = np.max(f.variables['gdept_0'][0,:,y,:]*f.variables['tmask'][0,:,y,:], axis=0)


# In[31]:


fig, ax = plt.subplots(1,1,figsize=(16,6))
t=1
for run in paths:
    ax.plot(meshes[run], label=run)
ax.plot(bathy[y,:],'--k',label = 'Actual bathymetry')
ax.legend()
ax.set_ylim([450,0])
ax.grid()
ax.set_title('T-depth of bottom cell ')
ax.set_ylabel('Depth (m)')
ax.set_xlabel('x (km)')


# Now look at the grid spacing

# In[32]:


def isolate_bottom(grid_spacing):
    """Determines the grid spacing of the t-cell at the bottom of the ocean. 
    Looks for first place where grid_sacing is 0."""
    
    grid= np.zeros(grid_spacing.shape[1])
    for i in np.arange(grid_spacing.shape[1]):
        zero_flag = False
        for k in np.arange(grid_spacing.shape[0]):
            if zero_flag:
                break
            elif grid_spacing[k,i]==0:
                grid[i] = grid_spacing[k-1,i]
                zero_flag = True
    return grid


# In[33]:


spacing_1d = {}
spacing={}
y=5
for run in paths:
    f = nc.Dataset(paths[run])
    grid_spacing = np.expand_dims(f.variables['e3t_1d'][0,:],1)*f.variables['tmask'][0,:,y,:]
    #find grid spacing of lowest - 1D case
    spacing_1d[run] = isolate_bottom(grid_spacing)
    #same thing for the partial steps
    grid_spacing = f.variables['e3t_0'][0,:,y,:]*f.variables['tmask'][0,:,y,:]
    spacing[run] = isolate_bottom(grid_spacing)


# In[34]:


fig, axs = plt.subplots(3,1,figsize=(16,10))
t=1
ax=axs
colors = ['b','r','g']
for run ,ax in zip(paths,axs):
    ax.plot(spacing_1d[run], label='no partial applied')


for run ,ax in zip(paths,axs):
    ax.plot(spacing[run],'--', label='partial applied')
    ax.grid()
    ax.set_ylabel('T cell bottom grid spacing (m)')
    ax.set_title('{} spacing without/with partial steps'.format(run))
    ax.legend(loc=0)
    ax.set_ylim([0,35])


# There is some pretty strong variability in the partial step grid spacing. Does that matter? That variability is lower in the dbl_tanh case. What are the numerical restrictions on the vertical grid cell size?

# Look at the depth of the w point at the bottom.

# In[35]:


def isolate_bottom_w(grid_mask, grid):
    """Determines the depth of the w-cell at the bottom of the ocean. 
    Looks for k when grid_mask[k]=0, returns non masked grid at that point. """
    
    bottom= np.zeros(grid_mask.shape[1])
    for i in np.arange(grid_mask.shape[1]):
        zero_flag = False
        for k in np.arange(1,grid_mask.shape[0]):
            if zero_flag:
                break
            elif grid_mask[k,i]==0:
                bottom[i] = grid[k,i] 
                zero_flag = True
    return bottom


# In[36]:


bottom_1d = {}
bottom={}
y=5
for run in paths:
    f = nc.Dataset(paths[run])
    #1d case
    grid = np.expand_dims(f.variables['gdepw_1d'][0,:],1)*np.ones(f.variables['tmask'][0,:,y,:].shape)
    grid_mask = np.expand_dims(f.variables['gdepw_1d'][0,:],1)*f.variables['tmask'][0,:,y,:]
    bottom_1d[run] = isolate_bottom_w(grid_mask, grid)
    #partial
    grid = f.variables['gdepw_0'][0,:,y,:]
    grid_mask = f.variables['gdepw_0'][0,:,y,:]*f.variables['tmask'][0,:,y,:]
    bottom[run] = isolate_bottom_w(grid_mask, grid)


# In[37]:


fig, axs = plt.subplots(3,1,figsize=(16,10))
for run, ax in zip(paths, axs):
    ax.plot(bottom_1d[run], label ='no partial applied')
    ax.plot(bottom[run], label='partial applied ')
    ax.plot(bathy[y,:],'--k',label = 'Actual bathymetry')
    ax.legend()
    ax.set_ylim([450,0])
    ax.grid()
    ax.set_title('{} W-depth of bottom cell'.format(run))
    ax.set_ylabel('Depth (m)')
    ax.set_xlabel('x (km)')


# with the partial steps, we are representing the bathymetry very well in all cases. This is on the w grid. But, even when accounting for partial steps, the T grid bottom isn't very smooth..

# In[ ]: