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

Image(filename='/ocean/imachuca/MEOPAR/Ariane/results/English_Bay_Spill/Oil_Spill1.jpg')

Image(filename='/ocean/imachuca/MEOPAR/Ariane/results/English_Bay_Spill/Oil_Spill2.png')

grid = nc.Dataset('../../NEMO-forcing/grid/bathy_meter_SalishSea2.nc')
day1U = nc.Dataset('/data/dlatorne/MEOPAR/SalishSea/nowcast/09apr15/SalishSea_1h_20150409_20150409_grid_U.nc','r')
day2U = nc.Dataset('/data/dlatorne/MEOPAR/SalishSea/nowcast/10apr15/SalishSea_1h_20150410_20150410_grid_U.nc','r')
day4U = nc.Dataset('/data/dlatorne/MEOPAR/SalishSea/nowcast/12apr15/SalishSea_1h_20150412_20150412_grid_U.nc','r')
day1V = nc.Dataset('/data/dlatorne/MEOPAR/SalishSea/nowcast/09apr15/SalishSea_1h_20150409_20150409_grid_V.nc','r')
ugrid = day1U.variables['vozocrtx']
vgrid = day1V.variables['vomecrty']
zlevels = day1U.variables['depthu']
timesteps = day1U.variables['time_counter']

t=0
timestamp1 = nc_tools.timestamp(day1U, t)
timestamp2 = nc_tools.timestamp(day2U, t)
timestamp4 = nc_tools.timestamp(day4U, t)

bathy, X, Y = tidetools.get_bathy_data(grid)
lat = 49.289227
lon = -123.187080
yind, xind = tidetools.find_closest_model_point(lon, lat, X, Y, bathy)
print xind, yind

zlevel = 0
x_slice = np.arange(260, 370)
y_slice = np.arange(420, 520)

ugrid_tzyx = np.ma.masked_values(ugrid[t, zlevel, y_slice, x_slice], 0)
vgrid_tzyx = np.ma.masked_values(vgrid[t, zlevel, y_slice, x_slice], 0)
u_tzyx, v_tzyx = viz_tools.unstagger(ugrid_tzyx, vgrid_tzyx)

fig, ax = plt.subplots(1, 1, figsize=(20, 10))
viz_tools.set_aspect(ax)

ax.streamplot(x_slice[1:], y_slice[1:], u_tzyx, v_tzyx, density=2)
ax.plot(xind, yind, marker='o', markersize=14, markeredgewidth=2, color='red', label='Origin')
viz_tools.plot_land_mask(ax, grid, xslice=x_slice, yslice=y_slice)

ax.set_xlim(x_slice[0], x_slice[-1])
ax.set_ylim(y_slice[0], y_slice[-1])

ax.set_xlabel('x Index')
ax.set_ylabel('y Index')
timestamp = nc_tools.timestamp(day1U, t)
ax.set_title('Velocity Field\n'+timestamp1.format('DD-MMM-YYYY, HH:mm [UTC]'))

legend = ax.legend(numpoints=1, loc=4)

part1 = nc.Dataset('/ocean/imachuca/MEOPAR/Ariane/results/English_Bay_Spill/0_4_8_12hrs/ariane_trajectories_qualitative.nc','r');
lon1=part1.variables['traj_lon']
lat1=part1.variables['traj_lat']
dep1=part1.variables['traj_depth']
tim1=part1.variables['traj_time']
inx1=part1.variables['init_x']
iny1=part1.variables['init_y']
int1=part1.variables['init_t']

part2 = nc.Dataset('/ocean/imachuca/MEOPAR/Ariane/results/English_Bay_Spill/12_16_20_12hrs/ariane_trajectories_qualitative.nc','r');
lon2=part2.variables['traj_lon']
lat2=part2.variables['traj_lat']
dep2=part2.variables['traj_depth']
tim2=part2.variables['traj_time']
inx2=part2.variables['init_x']
iny2=part2.variables['init_y']
int2=part2.variables['init_t']

part3 = nc.Dataset('/ocean/imachuca/MEOPAR/Ariane/results/English_Bay_Spill/0_48hrs/ariane_trajectories_qualitative.nc','r');
lon3=part3.variables['traj_lon']
lat3=part3.variables['traj_lat']
dep3=part3.variables['traj_depth']
tim3=part3.variables['traj_time']
inx3=part3.variables['init_x']
iny3=part3.variables['init_y']
int3=part3.variables['init_t']

partt = nc.Dataset('/ocean/imachuca/MEOPAR/Ariane/results/English_Bay_Spill/test_15min/ariane_trajectories_qualitative.nc','r');
lont=partt.variables['traj_lon']
latt=partt.variables['traj_lat']
dept=partt.variables['traj_depth']
timt=partt.variables['traj_time']
inxt=partt.variables['init_x']
inyt=partt.variables['init_y']
intt=partt.variables['init_t']

part4 = nc.Dataset('/ocean/imachuca/MEOPAR/Ariane/results/English_Bay_Spill/48hrs_multiples/ariane_trajectories_qualitative.nc','r');
lon4=part4.variables['traj_lon']
lat4=part4.variables['traj_lat']
dep4=part4.variables['traj_depth']
tim4=part4.variables['traj_time']
inx4=part4.variables['init_x']
iny4=part4.variables['init_y']
int4=part4.variables['init_t']

def plot_map(ax, grid):
    fig, ax = plt.subplots(1,1,figsize=(20,15))
    viz_tools.plot_land_mask(ax,grid,coords='map')
    viz_tools.set_aspect(ax)
    ax.set_xlim([-123.29,-123.13])
    ax.set_ylim([49.22,49.34])
    return ax

def add_labels(ax):
    ax.plot(-123.155, 49.274, marker='o', markersize=14, markeredgewidth=2, color='MistyRose') #Kitsilano Beach
    ax.plot(-123.15, 49.30, marker='o', markersize=14, markeredgewidth=2, color='MistyRose') #Third Beach
    ax.plot(-123.216, 49.277, marker='o', markersize=14, markeredgewidth=2, color='MistyRose') #Spanish Banks
    bbox_args = dict(boxstyle='square', facecolor='white')
    ax.annotate('Spanish\nBanks', (-123.22, 49.265), fontsize=15, color='white')
    ax.annotate('Kitsilano\nBeach', (-123.155, 49.26), fontsize=15, color='white')
    ax.annotate('Third\nBeach', (-123.145, 49.292), fontsize=15, color='white')
    return ax

ax = plot_map(ax, grid)
ax = add_labels(ax)
ax.set_title('Particles originating at point of oil spill in English Bay\n'+timestamp1.format('DD-MMM-YYYY')
             +', Duration = 12 hrs, Frequency = 15 min' )
legend.get_title().set_fontsize('20')
ax.annotate('Origin', (-123.19, 49.282), fontsize=20)

n = np.arange(lon1.shape[1]+lon2.shape[1])
colors = cm.rainbow(np.linspace(0, 1, len(n)))
for N, c in zip(n[0:lon1.shape[1]], colors[0:3]):
    ax.plot(lon1[1:,N],lat1[1:,N], linewidth=3.0, color=c, label=int1[N]-0.5)
    ax.plot(lon1[0,N],lat1[0,N], marker='o', markersize=14, markeredgewidth=2, color='MistyRose')
    
for N, c in zip(n[0:lon2.shape[1]], colors[3:7]):
    ax.plot(lon2[1:,N],lat2[1:,N], linewidth=3.0, color=c, label=int2[N]-0.5)
    ax.plot(lon2[0,N],lat2[0,N], marker='o', markersize=14, markeredgewidth=2, color='MistyRose')
    
ax.legend(loc=4, ncol=3, title=r'Initial Times of Particle Trajectories [UTC]')

ax = plot_map(ax, grid)
ax = add_labels(ax)
ax.set_title('Particles originating at point of oil spill in English Bay\n'+timestamp1.format('DD-MMM-YYYY')+' to '+timestamp2.format('DD-MMM-YYYY')
             +', Duration = 48 hrs, Frequency = 15 min' )

ax.plot(lon3[1:,0],lat3[1:,0], linewidth=3.0, color='red')
ax.plot(lon3[0,0],lat3[0,0], marker='o', markersize=14, markeredgewidth=2, color='MistyRose', label='Origin')
ax.annotate('Origin', (-123.19, 49.282), fontsize=20)

times = np.arange(0.5,48.75,0.25)
len_times = len(times)
times_ind = np.arange(len_times)
times_ind_r = times_ind[::-1]
times_ind_r

ax = plot_map(ax, grid)
ax = add_labels(ax)
ax.annotate('Origin', (-123.19, 49.282), fontsize=20)
ax.set_title('Particles originating at point of oil spill in English Bay\n'+
             'One particle dropped every 15 minutes from '
             +timestamp1.format('DD-MMM-YYYY') + ', 00:00'
             +' to '+timestamp2.format('DD-MMM-YYYY') + ', 24:00\n'
             +'Duration = Trajectories end on '+timestamp2.format('DD-MMM-YYYY') + ', 24:00'
             +', Frequency of particle positions = 15 min' )
for i, j in zip(times_ind, times_ind_r):
    ax.plot(lon4[0:j+1,i],lat4[0:j+1,i], color='DodgerBlue')
    ax.plot(lon4[0,0], lat4[0,0], marker='o', markersize=14, markeredgewidth=2, color='MistyRose')
plt.show()

ax = plot_map(ax, grid)
ax = add_labels(ax)
ax.annotate('Origin', (-123.19, 49.282), fontsize=20)
ax.set_title('Particles originating at point of oil spill in English Bay\n'+
             'One particle dropped every 15 minutes from '
             +timestamp1.format('DD-MMM-YYYY')+', 00:00'
             +' to '+timestamp2.format('DD-MMM-YYYY')+', 24:00\n'
             +'Duration = 48 hrs (Final trajectory ends on '+timestamp4.format('DD-MMM-YYYY')+', 24:00\n'
             +'Frequency of particle positions = 15 min' )
ax.plot(lon4[1:,:],lat4[1:,:], color='DodgerBlue')
ax.plot(lon4[0,0], lat4[0,0], marker='o', markersize=14, markeredgewidth=2, color='MistyRose')
plt.show()

ax = plot_map(ax, grid)
ax = add_labels(ax)
ax.annotate('Origin', (-123.19, 49.282), fontsize=20)
legend.get_title().set_fontsize('20')
ax.set_title('Particles originating at point of oil spill in English Bay\n'+
             'One particle dropped every 15 minutes from '
             +timestamp1.format('DD-MMM-YYYY')+', 00:00'
             +' to '+timestamp2.format('DD-MMM-YYYY')+', 24:00\n'
             +'Duration = 48 hrs (Final trajectory ends on '+timestamp4.format('DD-MMM-YYYY')+', 24:00\n'
             +'Frequency of particle positions = 15 min' )
ax.plot(lon4[1:, 0],lat4[1:, 0], color='MediumSlateBlue', label='09-Apr')
ax.plot(lon4[1:, 0:96],lat4[1:, 0:96], color='MediumSlateBlue')
ax.plot(lon4[1:, 96],lat4[1:, 96], color='DodgerBlue', label='10-Apr')
ax.plot(lon4[1:, 96:],lat4[1:, 96:], color='DodgerBlue')
ax.plot(lon4[0,0], lat4[0,0], marker='o', markersize=14, markeredgewidth=2, color='MistyRose')
ax.legend(loc=4, ncol=1, title=r'Release date')
plt.show()