import netCDF4 as NC
import numpy as np

import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
import matplotlib.cm as cm
import datetime as dt

from salishsea_tools import (nc_tools,viz_tools,tidetools,stormtools)

%matplotlib inline

grid = NC.Dataset('/ocean/imachuca/MEOPAR/NEMO-forcing/grid/bathy_meter_SalishSea2.nc','r')
bathy, X, Y = tidetools.get_bathy_data(grid)


lats = grid.variables['nav_lat']
lons = grid.variables['nav_lon']
bath = grid.variables['Bathymetry']

fig, ax = plt.subplots(1, 1, figsize=(20, 10))
viz_tools.set_aspect(ax)
cmap = plt.get_cmap('Blues')
cmap.set_bad('burlywood')
mesh = ax.pcolormesh(bathy, cmap=cmap)
fig.colorbar(mesh)
plt.axis((30, 200, 670, 900))

xx = [120,91]
yy = [744,855]
xx2 = [135,156,182,140,142,125,126,130,149,119,84,53]
yy2 = [680,680,680,710,740,762,786,807,821,836,858,890]

labb = ['Campbell River','Johnstone Strait']
labb2 = np.arange(1,len(xx2)+1)

for x,y,lab in zip (xx,yy,labb):
    ax.scatter(x,y,color='white',marker='s')
    ax.annotate(lab,(x,y),fontsize=10,color='white')
for x2,y2,lab2 in zip (xx2,yy2,labb2):
    ax.scatter(x2,y2,color='red',marker='.')
    ax.annotate(lab2,(x2,y2),fontsize=15,color='black')

out = NC.Dataset('/ocean/imachuca/MEOPAR/Ariane/results/johnstone/out/ariane_trajectories_qualitative.nc','r')
lon1=out.variables['traj_lon']
lat1=out.variables['traj_lat']
dep1=out.variables['traj_depth']
x1=out.variables['init_x']
y1=out.variables['init_y']
t1=out.variables['traj_time']

iin = NC.Dataset('/ocean/imachuca/MEOPAR/Ariane/results/johnstone/in/ariane_trajectories_qualitative.nc','r')
lon2=iin.variables['traj_lon']
lat2=iin.variables['traj_lat']
dep2=iin.variables['traj_depth']
x2=iin.variables['init_x']
y2=iin.variables['init_y']
t2=iin.variables['traj_time']

len(x1),len(x2)

cycle = NC.Dataset('/ocean/imachuca/MEOPAR/Ariane/results/johnstone/cycle/ariane_trajectories_qualitative.nc','r')
lon_cyc=cycle.variables['traj_lon']
lat_cyc=cycle.variables['traj_lat']
dep_cyc=cycle.variables['traj_depth']
x_cyc=cycle.variables['init_x']
y_cyc=cycle.variables['init_y']
t_cyc=cycle.variables['traj_time']
ti_cyc=cycle.variables['init_t']
lon_cyc.shape

one = NC.Dataset('/ocean/imachuca/MEOPAR/Ariane/results/johnstone/one/ariane_trajectories_qualitative.nc','r')
lon_one=one.variables['traj_lon']
lat_one=one.variables['traj_lat']
dep_one=one.variables['traj_depth']
x_one=one.variables['init_x']
y_one=one.variables['init_y']
t_one=one.variables['traj_time']
ti_one=one.variables['init_t']
lon_one.shape

fig, (ax1,ax2) = plt.subplots(1,2,figsize=(20,10))
land_colour = 'burlywood'

cmap = plt.get_cmap('Blues')
mesh = ax1.pcolormesh(lons[:], lats[:], bathy[:], cmap=cmap)
mesh = ax2.pcolormesh(lons[:], lats[:], bathy[:], cmap=cmap)

colors = cm.gist_rainbow(np.linspace(0, 1, len(x1)))
n=np.arange(len(x1))
labb = np.arange(1,len(x1)+1)

for c,N,lab in zip(colors,n,labb):
    ax1.scatter(lon1[1:,N],lat1[1:,N],color=c,marker='.')
    ax1.scatter(lon1[0,N],lat1[0,N],color='0.30',marker='s')
    ax1.annotate(lab,(lon1[0,N],lat1[0,N]),fontsize=15,color='black')
    
colors = cm.gist_rainbow(np.linspace(0, 1, len(x2)))
n=np.arange(len(x2))
labb = np.arange(8,len(x2)+8)
    
for c,N,lab in zip(colors,n,labb):
    ax2.scatter(lon2[1:,N],lat2[1:,N],color=c,marker='.')
    ax2.scatter(lon2[0,N],lat2[0,N],color='0.30',marker='s')
    ax2.annotate(lab,(lon2[0,N],lat2[0,N]),fontsize=15,color='black')

viz_tools.plot_land_mask(ax1,grid,coords='map', color=land_colour)
viz_tools.plot_land_mask(ax2,grid,coords='map', color=land_colour)

ax1.set_xlim([-125.5,-124.55])
ax1.set_ylim([49.65,50.37])

ax2.set_xlim([-126.2,-125.3])
ax2.set_ylim([50.3,50.5])


ax1.set_title('2 days, 30 sec frequency, time 0, Surface')
ax1.set_xlabel('lon')
ax1.set_ylabel('lat')

ax2.set_title('2 days, 30 sec frequency, time 0, Surface')
ax2.set_xlabel('lon')
ax2.set_ylabel('lat')

wlev = stormtools.load_observations('13-Dec-2006', '16-Dec-2006','CampbellRiver')
fig,ax=plt.subplots(1,1,figsize=(20,5))
ax.plot(wlev.time,wlev.slev,'-k')
ax.set_xlim(dt.datetime(2006,12,14),dt.datetime(2006,12,16))
ax.set_ylabel('Water level (m)')
ax.set_xlabel('Dec 14 - Dec 15(hrs)')
ax.set_title('Observed water level at Campbell River')

ax.plot(wlev.time[17],wlev.slev[17],'o',color='Maroon',markersize=9)
ax.plot(wlev.time[19],wlev.slev[19],'o',color='MediumBlue',markersize=9)
ax.plot(wlev.time[22],wlev.slev[22],'o',color='Olive',markersize=9)
ax.plot(wlev.time[25],wlev.slev[25],'o',color='DarkViolet',markersize=9)
ax.plot(wlev.time[30],wlev.slev[30],'o',color='Tomato',markersize=9)
ax.plot(wlev.time[33],wlev.slev[33],'o',color='DeepSkyBlue',markersize=9)
ax.plot(wlev.time[36],wlev.slev[36],'o',color='SpringGreen',markersize=9)
ax.plot(wlev.time[39],wlev.slev[39],'o',color='DeepPink',markersize=9)

print wlev.time[17]
print wlev.time[19]
print wlev.time[22]
print wlev.time[25]
print wlev.time[30]
print wlev.time[33]
print wlev.time[36]
print wlev.time[39]

(fig,axs)=plt.subplots(8,2,figsize=(20,50))
times = 8
particles=4

n=np.arange(times)
tt = [17,19,22,25,30,33,36,39]
colorss = ['Maroon','MediumBlue','Olive','DarkViolet','Tomato','DeepSkyBlue','SpringGreen','DeepPink']

for ax,t,colors,N in zip(axs[:,1],tt,colorss,n):
    ax.plot(wlev.time[:],wlev.slev[:],'-k')
    ax.set_xlim(dt.datetime(2006,12,14),dt.datetime(2006,12,15))
    ax.set_xticklabels([])
    ax.set_ylabel('Water level [m]]')
    ax.set_xlabel('Dec 14, 00:00 - 24:00')
    ax.plot(wlev.time[t],wlev.slev[t],'o',color=colors,markersize=9)
    ax.set_title(str(wlev.time[t]))
    ax.set_position((0.83, (N+1)*0.125, 0.2, 0.1)) 


n=np.arange(times)
m=np.arange(particles)
for ax,N in zip(axs[:,0], n):
    for M,colors in zip(m,colorss):
        A = (particles*N)+M
        ax.scatter(lon_cyc[0,A],lat_cyc[0,A],color='0.30',marker='s')
        ax.scatter(lon_cyc[1:,A],lat_cyc[1:,A],color=colors,marker='o')
    viz_tools.plot_land_mask(ax,grid,coords='map', color='burlywood')
    ax.set_xlim([-125.55,-125.1])
    ax.set_ylim([49.93,50.37])
    ax.set_title(str(ti_cyc[N*particles]) + 'hrs')
    ax.set_position((0.10, (N+1)*0.125, 0.7, 0.116))

(fig,ax)=plt.subplots(1,1,figsize=(10,8))
viz_tools.plot_land_mask(ax,grid,coords='map', color='burlywood')
ax.set_xlim([-125.55,-125.2])
ax.set_ylim([50,50.3])
ax.scatter(lon_one[0,:],lat_one[0,:],color='0.30',marker='s')
ax.scatter(lon_one[1:,:],lat_one[1:,:],color='Olive',marker='.')