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

# Compare stratification and mixing under different vertical resolution

# In[2]:


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[3]:


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

run1 = 'dbl_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])


# In[4]:


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[5]:


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[6]:


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)


# * After 40 days, water that has piled over the sill is saltier.
# * Also, more salty water is entering the basin.
# * This suggests less mixing over the sills.
# * I haven't looked closely at the surface because I don't think we want the resolution up there. Will look at Susan's next.

# In[10]:


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, np.arange(5,15.1,.1),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)


# * Colder water into the basin

# In[11]:


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(-2,2.1,.1),ls='None')
        ax.set_title('U [m/s]: t={} : {}'.format(t, run))
        plt.colorbar(mesh,ax=ax)


# 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):
        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[13]:


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)


# * High diffusivity at at the bottom of the basin. I think the plunging gravity current sets up a seiche like motion in the basin. Too much salty water is entering and this is partially mixed out due to motions in the basin. 
# * Hard to tell by eye if the mixing coefficients are high in one than the other... Need to average somehow.
# * By eye, looks like more mixing in surface layer in the new resolution.

# 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):
        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[15]:


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


# In[18]:


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[19]:


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

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


# * On average, over the mixing area, the newer resolultion has higher eddy coefficients. Is that because there is more mixing in the surface layer? We might fix this by using Susan's scheme...
# * Can I isolate the surface layer?

# ###Surface

# In[20]:


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

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


# ###Intermediate

# In[21]:


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

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


# ###Deep

# In[22]:


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


# * It looks like eddy coefficients are, on average, higher at all depths. Why? Is that inconsistent with saltier water over sills?
# * But, the max depth averaged diffusivity in the region of interest is consistently higher in the old resolution case. 

# ###Strait of Juan de Fuca

# In[23]:


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 ---> Less mixing before entering sill area?

# ### Strait of Georgia

# In[24]:


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


# ### Whole domain

# In[25]:


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


# # Stratification Profiles

# In[26]:


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[27]:


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[28]:


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


# Slightly fresher surface layer with the new resolution.

# In[ ]: