from __future__ import division
import matplotlib.pyplot as plt
import netCDF4 as nc
import numpy as np
from salishsea_tools import (nc_tools,viz_tools)
from IPython.core.display import Image
import matplotlib.cm as cm
from mpl_toolkits.mplot3d import Axes3D
from salishsea_tools import tidetools
%matplotlib inline

tracers = nc.Dataset('/ocean/dlatorne/MEOPAR/SalishSea/results/spin-up/18sep27sep/SalishSea_1d_20030918_20030927_grid_T.nc')
sal = tracers.variables['vosaline']
zlevels = tracers.variables['deptht']
ssal = sal.shape
ssal

tracersW = nc.Dataset('/ocean/dlatorne/MEOPAR/SalishSea/results/spin-up/18sep27sep/SalishSea_1d_20030918_20030927_grid_W.nc')
nc_tools.show_variables(tracers)
velW=tracersW.variables['vovecrtz']
svelW = velW.shape
svelW

grid = nc.Dataset('/ocean/imachuca/MEOPAR/NEMO-forcing/grid/bathy_meter_SalishSea2.nc','r')
nc_tools.show_variables(grid)
lats = grid.variables['nav_lat']
lons = grid.variables['nav_lon']
bathy = grid.variables['Bathymetry']
bathy.shape
bath, X, Y = tidetools.get_bathy_data(grid)

!head thalweg.txt
thalweg = np.loadtxt('thalweg.txt', dtype=int, unpack=True)

PT_thal = nc.Dataset('/ocean/imachuca/MEOPAR/Ariane/results/thalweg_islands/along_thalweg/unconfined/ariane_trajectories_qualitative.nc','r')
lonT=PT_thal.variables['traj_lon']
latT=PT_thal.variables['traj_lat']
depT=PT_thal.variables['traj_depth']
xT=PT_thal.variables['init_x']
yT=PT_thal.variables['init_y']
tT=PT_thal.variables['traj_time']
nc_tools.show_variables(PT_thal)

PT_isl = nc.Dataset('/ocean/imachuca/MEOPAR/Ariane/results/thalweg_islands/islands/unconfined/ariane_trajectories_qualitative.nc','r')
lonI=PT_isl.variables['traj_lon']
latI=PT_isl.variables['traj_lat']
depI=PT_isl.variables['traj_depth']
xI=PT_isl.variables['init_x']
yI=PT_isl.variables['init_y']
tI=PT_thal.variables['traj_time']

PTN_thal = nc.Dataset('/ocean/imachuca/MEOPAR/Ariane/results/thalweg_islands/along_thalweg/nointerp/ariane_trajectories_qualitative.nc','r')
lonTN=PTN_thal.variables['traj_lon']
latTN=PTN_thal.variables['traj_lat']
depTN=PTN_thal.variables['traj_depth']
xTN=PTN_thal.variables['init_x']
yTN=PTN_thal.variables['init_y']
tTN=PTN_thal.variables['traj_time']
nc_tools.show_variables(PTN_thal), latTN.shape

PTN_isl = nc.Dataset('/ocean/imachuca/MEOPAR/Ariane/results/thalweg_islands/islands/nointerp/ariane_trajectories_qualitative.nc','r')
lonIN=PTN_isl.variables['traj_lon']
latIN=PTN_isl.variables['traj_lat']
depIN=PTN_isl.variables['traj_depth']
xIN=PTN_isl.variables['init_x']
yIN=PTN_isl.variables['init_y']
tIN=PTN_thal.variables['traj_time']
nc_tools.show_variables(PTN_isl), latIN.shape

Image(filename='ThalSal.png')

Image(filename='ThalDepth.png') 

Image(filename='ThalMap.png') 

# Set up the figure and axes
fig, axr= plt.subplots(1,1, figsize=(30, 15))
land_colour = 'indigo'
axr.set_axis_bgcolor(land_colour)
axr.set_position((0.83, 0.125, 0.2, 0.775))

# Plot thalweg points on bathymetry map
viz_tools.set_aspect(axr)
cmap = plt.get_cmap('winter_r')
cmap.set_bad(land_colour)
bathy = grid.variables['Bathymetry']
x_slice = np.arange(bathy.shape[1])
y_slice = np.arange(200, 800)
axr.pcolormesh(x_slice, y_slice, bathy[y_slice, x_slice], cmap=cmap)
axr.plot(thalweg[1], thalweg[0], linestyle='-', marker='+', color='wheat', label='Thalweg Points')
legend = axr.legend(loc='best', fancybox=True, framealpha=0.25)
axr.set_xlabel('x Index')
axr.set_ylabel('y Index')
axr.grid()
axr.set_xlim([200,355])
axr.set_ylim([240,450])

xk = [268,272,285,300,289,262,239]
yk = [397,385,365,350,342,345,330]
xw = [320,293,266,324,292,265,302]
yw = [350,331,317,316,295,287,274]
lab = np.arange(1,8)

for xxk, yyk, xxw, yyw, labb in zip (xk,yk,xw,yw,lab):
    axr.scatter(xxk,yyk,color='black',marker='s')
    axr.scatter(xxw,yyw,color='white',marker='s')
    axr.annotate(labb,(xxk,yyk),fontsize=15,color='black')
    axr.annotate(labb,(xxw,yyw),fontsize=15,color='white')


bathyk = np.zeros(7)
bathyw = np.zeros(7)

for i in range (7):
    bathyk[i] = 0.5*bathy[yk[i],xk[i]]
    bathyw[i] = 0.5*bathy[yw[i],xw[i]]
bathyk,bathyw

zlevels[:]

iyTN = np.zeros(shape=(10,7))
ixTN = np.zeros(shape=(10,7))
iyIN = np.zeros(shape=(10,7))
ixIN = np.zeros(shape=(10,7))
depthTNS=np.empty_like(iyTN)
depthINS=np.empty_like(iyIN)

partcc = np.arange(7)
dyss = np.arange(10)

for partc in zip(partcc):
    for dys in zip(dyss):
        iyTN,ixTN=  tidetools.find_closest_model_point(lonTN[dys,partc], latTN[dys,partc], X, Y, bath, allow_land=False)
        iyIN, ixIN = tidetools.find_closest_model_point(round(lonIN[dys,partc],4), round(latIN[dys,partc],4), X, Y, bath, allow_land=False)
        depthTNS[dys,partc]=bathy[int(float(iyTN)),int(float(ixTN))]
        depthINS[dys,partc]=bathy[int(float(iyIN)),int(float(ixIN))]

(fig,axs)=plt.subplots(7,2,figsize=(15,70))
n = np.arange(7)

for ax,N in zip(axs[:,0],n):
    aspect = viz_tools.set_aspect(ax, coords='map', lats=lats)
    cmap = plt.get_cmap('Blues')
    cmap.set_bad('burlywood')
    mesh = ax.pcolormesh(lons[:], lats[:], bathy[:], cmap=cmap)
    ax.set_xlim([-123.8,-122.7])
    ax.set_ylim([48,49.6])
    if (ax==axs[6,0]):
        ax.set_xlim([-124.6,-122.7])
        ax.set_ylim([48,49.6])
    ax.scatter(lonT[0,N],latT[0,N],color='red',marker='s', label='Initial Position')
    cm = plt.cm.get_cmap('RdYlBu')
    ax.scatter(lonT[1:,N],latT[1:,N],c=tT, cmap=cm,marker='o', label='Interp. Trajectory', alpha=0.7)
    ax.set_xlabel('Longitude')
    ax.set_ylabel('Latitude')
    ax.set_title('Trajectory for particle along thalweg')
    ax.legend(scatterpoints = 1, loc='lower right')

for ax,N in zip(axs[:,1],n):
    cm = plt.cm.get_cmap('RdYlBu')
    ax.plot(tT[:,N]*10, depT[:,N]*-1, '-.', label = 'Interpolated', color='grey')
    ax.plot(tTN[:,N]*10, depTN[:,N]*-1, '-o', label = 'Not Interpolated', color='green')
    ax.plot(tTN[:,N]*10,depthTNS[:,N],'-o', color='blue', label='Bathymetry')
    ax.set_ylim(400, 0)
    ax.set_xticks(np.arange(0, 10, 1))
    ax.set_xticklabels(np.arange(0, 10, 1))
    ax.set_xlabel('Time [days]')
    ax.set_ylabel('Depth [m]')
    ax.set_title('Depth of particle along trajectory')
    ax.legend(scatterpoints = 1, loc='lower left')


(fig,axs)=plt.subplots(7,2,figsize=(15,70))
n = np.arange(7)

for ax,N in zip(axs[:,0],n):
    aspect = viz_tools.set_aspect(ax, coords='map', lats=lats)
    cmap = plt.get_cmap('Blues')
    cmap.set_bad('burlywood')
    mesh = ax.pcolormesh(lons[:], lats[:], bathy[:], cmap=cmap)
    ax.set_xlim([-123.5,-122.5])
    ax.set_ylim([48.2,49])
    ax.scatter(lonIN[0,N],latIN[0,N],color='red',marker='s', label='Initial Position')
    cm = plt.cm.get_cmap('RdYlBu')
    ax.scatter(lonI[1:,N],latI[1:,N],c=tI, cmap=cm,marker='o', label='Interp. Trajectory', alpha=0.7)
    ax.set_xlabel('Longitude')
    ax.set_ylabel('Latitude')
    ax.set_title('Trajectory for particle around islands')
    ax.legend(scatterpoints = 1, loc='lower left')

for ax,N in zip(axs[:,1],n):
    cm = plt.cm.get_cmap('RdYlBu')
    ax.plot(tI[:,N]*10, depI[:,N]*-1, '-.', label = 'Interpolated', color='grey')
    ax.plot(tIN[:,N]*10, depIN[:,N]*-1, '-o', label = 'Not Interpolated', color='green')
    ax.plot(tIN[:,N]*10,depthINS[:,N],'-o', color='blue', label='Bathymetry')
    ax.set_ylim(400, 0)
    ax.set_xticks(np.arange(0, 10, 1))
    ax.set_xticklabels(np.arange(0, 10, 1))
    ax.set_xlabel('Time [days]')
    ax.set_ylabel('Depth [m]')
    ax.set_title('Depth of particle along trajectory')
    ax.legend(scatterpoints = 1, loc='lower left')