%matplotlib inline
from matplotlib import pylab
import matplotlib.pyplot as plt
import netCDF4 as NC
import numpy as np
from mpl_toolkits.mplot3d import Axes3D
from salishsea_tools import tidetools

from salishsea_tools import (nc_tools,viz_tools)

import matplotlib.cm as cm

from matplotlib import animation

from numpy import *
from pylab import *

import Image

test2 = NC.Dataset('/ocean/imachuca/MEOPAR/Ariane/results/riverparticles/times/test2/ariane_trajectories_qualitative.nc','r')

test3 = NC.Dataset('/ocean/imachuca/MEOPAR/Ariane/results/riverparticles/times/test3/ariane_trajectories_qualitative.nc','r')

test4 = NC.Dataset('/ocean/imachuca/MEOPAR/Ariane/results/riverparticles/times/test4/ariane_trajectories_qualitative.nc','r')

test5 = NC.Dataset('/ocean/imachuca/MEOPAR/Ariane/results/riverparticles/times/test5/ariane_trajectories_qualitative.nc','r')

test6 = NC.Dataset('/ocean/imachuca/MEOPAR/Ariane/results/riverparticles/times/test6/ariane_trajectories_qualitative.nc','r')

test7 = NC.Dataset('/ocean/imachuca/MEOPAR/Ariane/results/riverparticles/times/test7/ariane_trajectories_qualitative.nc','r')

test8 = NC.Dataset('/ocean/imachuca/MEOPAR/Ariane/results/riverparticles/times/test8/ariane_trajectories_qualitative.nc','r')

fordocs = NC.Dataset('/ocean/imachuca/MEOPAR/Ariane/results/fordocs/ariane_trajectories_qualitative.nc','r')

lon2=test2.variables['traj_lon']
lat2=test2.variables['traj_lat']
dep2=test2.variables['traj_depth']
xs2=test2.variables['init_x']
ys2=test2.variables['init_y']

lon3=test3.variables['traj_lon']
lat3=test3.variables['traj_lat']
dep3=test3.variables['traj_depth']
xs3=test3.variables['init_x']
ys3=test3.variables['init_y']

lon4=test4.variables['traj_lon']
lat4=test4.variables['traj_lat']
dep4=test4.variables['traj_depth']
xs4=test4.variables['init_x']
ys4=test4.variables['init_y']

lon5=test5.variables['traj_lon']
lat5=test5.variables['traj_lat']
dep5=test5.variables['traj_depth']
xs5=test5.variables['init_x']
ys5=test5.variables['init_y']

lon6=test6.variables['traj_lon']
lat6=test6.variables['traj_lat']
dep6=test6.variables['traj_depth']
xs6=test6.variables['init_x']
ys6=test6.variables['init_y']

lon7=test7.variables['traj_lon']
lat7=test7.variables['traj_lat']
dep7=test7.variables['traj_depth']
xs7=test7.variables['init_x']
ys7=test7.variables['init_y']

lon8=test8.variables['traj_lon']
lat8=test8.variables['traj_lat']
dep8=test8.variables['traj_depth']
xs8=test8.variables['init_x']
ys8=test8.variables['init_y']

lond=fordocs.variables['traj_lon']
latd=fordocs.variables['traj_lat']
depd=fordocs.variables['traj_depth']
xsd=fordocs.variables['init_x']
ysd=fordocs.variables['init_y']

grid = NC.Dataset('/ocean/imachuca/MEOPAR/NEMO-forcing/grid/bathy_meter_SalishSea2.nc','r')

fig, ax = plt.subplots(1,1,figsize=(10,6.6))
ax.scatter(lon2[1:],lat2[1:],c='purple')
ax.scatter(lon2[0],lat2[0],color='0.30',marker='s')
ax.set_title('Test 2: time = 1')
viz_tools.plot_land_mask(ax,grid,coords='map')
ax.set_xlim([-123.5,-122.9])
ax.set_ylim([48.9,49.45])

fig, ax = plt.subplots(1,1,figsize=(10,6.6))
ax.scatter(lon3[1:],lat3[1:],c='purple')
ax.scatter(lon3[0],lat3[0],color='0.30',marker='s')
ax.set_title('Test 3: time = 24, nbout = 47, lmt = 48')
viz_tools.plot_land_mask(ax,grid,coords='map')
ax.set_xlim([-123.5,-122.9])
ax.set_ylim([48.9,49.45])

fig, ax = plt.subplots(1,1,figsize=(10,6.6))
ax.scatter(lon4[1:],lat4[1:],c='purple')
ax.scatter(lon4[0],lat4[0],color='0.30',marker='s')
ax.set_title('Test 4: time = 24, nbout = 24, lmt = 48')
viz_tools.plot_land_mask(ax,grid,coords='map')
ax.set_xlim([-123.5,-122.9])
ax.set_ylim([48.9,49.45])

fig, ax = plt.subplots(1,1,figsize=(10,6.6))
ax.scatter(lon5[1:],lat5[1:],c='purple')
ax.scatter(lon5[0],lat5[0],color='0.30',marker='s')
ax.set_title('Test 5: time = 24, nbout = 47, lmt = 48, ntfic=1')
viz_tools.plot_land_mask(ax,grid,coords='map')
ax.set_xlim([-123.5,-122.9])
ax.set_ylim([48.9,49.45])

fig, ax = plt.subplots(1,1,figsize=(10,6.6))
ax.scatter(lon6[1:],lat6[1:],c='purple')
ax.scatter(lon6[0],lat6[0],color='0.30',marker='s')
ax.set_title('Test 6: time = 24, nbout = 24, lmt = 48, ntfic=1')
viz_tools.plot_land_mask(ax,grid,coords='map')
ax.set_xlim([-123.5,-122.9])
ax.set_ylim([48.9,49.45])

fig, ax = plt.subplots(1,1,figsize=(10,6.6))
ax.scatter(lon7[1:],lat7[1:],c='purple')
ax.scatter(lon7[0],lat7[0],color='0.30',marker='s')
ax.set_title('Test 7: time = 24, nbout = 25, lmt = 48, ntfic=1')
viz_tools.plot_land_mask(ax,grid,coords='map')
ax.set_xlim([-123.5,-122.9])
ax.set_ylim([48.9,49.45])

fig, ax = plt.subplots(1,1,figsize=(10,6.6))
ax.scatter(lon8[1:],lat8[1:],c='purple')
ax.scatter(lon8[0],lat8[0],color='0.30',marker='s')
ax.set_title('Test 8: time = 1, nbout = 49, lmt = 48, ntfic=1')
viz_tools.plot_land_mask(ax,grid,coords='map')
ax.set_xlim([-123.5,-122.9])
ax.set_ylim([48.9,49.45])

n = np.arange(5)
colors = cm.rainbow(np.linspace(0, 1, len(n)))
labss = ['depth 1, time 0.5', 'depth 1.5, time 0.5', 'depth 2, time 1.5', 'depth 1, time 0.5', 'depth 1, time 0.5']

fig, ax = plt.subplots(1,1,figsize=(15,10))

for N,c,labs in zip(n,colors,labss):
    ax.scatter(lond[1:,N],latd[1:,N],color=c,label=labs)
    ax.scatter(lond[0,N],latd[0,N],color='0.30',marker='s')

viz_tools.plot_land_mask(ax,grid,coords='map')
ax.set_xlim([-123.6,-122.7])
ax.set_ylim([48.5,49.4])
ax.legend()

ax.set_title('Confined particle trajectories')
ax.set_xlabel('Longitude')
ax.set_ylabel('Latitude')


fig = plt.figure(figsize=(20,10))
ax = fig.gca(projection='3d')

for N,c in zip(n,colors):
    ax.scatter(lond[1:,N],latd[1:,N],depd[1:,N],color=c)
    ax.scatter(lond[0,N],latd[0,N],depd[0,N],color='0.30',marker='s')
    
    
ax.set_xlabel('longitude')
ax.set_ylabel('latitude')
ax.set_zlabel('depth')
ax.set_title('Particle trajectories in 3D')

ax.view_init(elev=0, azim=-90)