%matplotlib inline
from matplotlib import pylab
import matplotlib.pyplot as plt
import netCDF4 as NC
import numpy as np

from salishsea_tools import (nc_tools,viz_tools)

# Load the data. Path name can be changed to look at different data.  
path = '/data/nsoontie/MEOPAR/SalishSea/results/BBL/eddy_noBBL/SalishSea_1d_20030819_20030828_grid_T.nc'
f = NC.Dataset(path,'r');

sal=f.variables['vosaline']
depths = f.variables['deptht'] #array of depths

T_lat = f.variables['nav_lat']
T_lon = f.variables['nav_lon']


smin=30; smax=32

lines = np.loadtxt('/data/nsoontie/MEOPAR/tools/analysis_tools/thalweg.txt', delimiter=" ", unpack=False)
lines = lines.astype(int)

thalweg_lon = T_lon[lines[:,0],lines[:,1]]
thalweg_lat = T_lat[lines[:,0],lines[:,1]]

ds=np.arange(0,lines.shape[0],1);
vs=np.arange(34,27.5,-0.5);
XX,ZZ = np.meshgrid(ds,-depths[:])



(fig,axs)=plt.subplots(10,1,figsize=(10,15))
ts=np.arange(10)

for t,ax in zip(ts,axs):
    salP=sal[t,:,lines[:,0],lines[:,1]];
    #masking
    mu =salP == 0; salP= np.ma.array(salP,mask=mu)
    mesh=ax.pcolormesh(XX,ZZ,salP,vmin=smin,vmax=smax,cmap='rainbow')
    fig.colorbar(mesh,ax=ax)
    CS=ax.contour(XX,ZZ,salP,vs, colors='black')
    ax.clabel(CS,fontsize=9, inline=1)
    ax.set_title('t = ' +str(t))


# Load the data. Path name can be changed to look at different data.  
path = '/data/nsoontie/MEOPAR/SalishSea/results/BBL/SalishSea_1d_20030819_20030828_grid_T.nc'
f = NC.Dataset(path,'r');

sal2=f.variables['vosaline']
depths = f.variables['deptht'] #array of depths

T_lat = f.variables['nav_lat']
T_lon = f.variables['nav_lon']


(fig,axs)=plt.subplots(10,2,figsize=(20,15))
ts=np.arange(10)

for t,ax in zip(ts,axs[:,0]):
    salP=sal[t,:,lines[:,0],lines[:,1]];
    #masking
    mu =salP == 0; salP= np.ma.array(salP,mask=mu)
    mesh=ax.pcolormesh(XX,ZZ,salP,vmin=smin,vmax=smax,cmap='rainbow')
    fig.colorbar(mesh,ax=ax)
    CS=ax.contour(XX,ZZ,salP,vs, colors='black')
    ax.clabel(CS,fontsize=9, inline=1)
    ax.set_title('t = ' +str(t))

for t,ax in zip(ts,axs[:,1]):
    salP2=sal2[t,:,lines[:,0],lines[:,1]];
    #masking
    mu =salP2 == 0; salP2= np.ma.array(salP2,mask=mu)
    mesh=ax.pcolormesh(XX,ZZ,salP2,vmin=smin,vmax=smax,cmap='rainbow')
    fig.colorbar(mesh,ax=ax)
    CS=ax.contour(XX,ZZ,salP2,vs, colors='black')
    ax.clabel(CS,fontsize=9, inline=1)
    ax.set_title('t = ' +str(t))

print 'Left: no  BBL'
print 'Right: BBL'

t=9
fig,axs=plt.subplots(3,1,figsize=(10,8))
smin=28; smax=31

ax=axs[0]
salP=sal[t,:,lines[:,0],lines[:,1]];
mu =salP == 0; salP= np.ma.array(salP,mask=mu)
mesh=ax.pcolormesh(XX,ZZ,salP,vmin=smin,vmax=smax,cmap='rainbow')
fig.colorbar(mesh,ax=ax)
CS=ax.contour(XX,ZZ,salP,vs, colors='black')
ax.clabel(CS,fontsize=9, inline=1)
ax.set_title('No BBL')
ax.set_xlim([0,700]); ax.set_ylim([-428,0])

ax=axs[1]
salP2=sal2[t,:,lines[:,0],lines[:,1]];
mu =salP2 == 0; salP2= np.ma.array(salP2,mask=mu)
mesh=ax.pcolormesh(XX,ZZ,salP2,vmin=smin,vmax=smax,cmap='rainbow')
fig.colorbar(mesh,ax=ax)
CS=ax.contour(XX,ZZ,salP2,vs, colors='black')
ax.clabel(CS,fontsize=9, inline=1)
ax.set_title('BBL')
ax.set_xlim([0,700]); ax.set_ylim([-428,0])

ax=axs[2]
mesh=ax.pcolormesh(XX,ZZ,salP2-salP,vmin=-.1,vmax=.1,cmap='jet')
ax.set_title('Difference')
ax.set_xlim([0,700]); ax.set_ylim([-428,0])
fig.colorbar(mesh,ax=ax)