from numpy import *
from pylab import *
import netCDF4 as NC
import numpy as np
import datetime as dt

from matplotlib import pylab
import matplotlib.pyplot as plt
import matplotlib.cm as cm
from matplotlib import animation

from salishsea_tools import (nc_tools,tidetools,viz_tools,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']

rd = NC.Dataset('/ocean/imachuca/MEOPAR/Ariane/results/riverparticles/depths/ariane_trajectories_qualitative.nc','r');
rd_lont=rd.variables['traj_lon']
rd_latt=rd.variables['traj_lat']
rd_dept=rd.variables['traj_depth']
rd_xs=rd.variables['init_x']
rd_ys=rd.variables['init_y']
rd_lont.shape

tracers = NC.Dataset('/data/nsoontie/MEOPAR/SalishSea/results/storm-surges/tide_fix/dec2006/all_forcing/1hour/SalishSea_1h_20061214_20061215_grid_T.nc')
sal = tracers.variables['vosaline']
ssh = tracers.variables['sossheig']
timesteps = tracers.variables['time_counter']
s_depth = tracers.variables['deptht']
s_lon = tracers.variables['nav_lon']
s_lat = tracers.variables['nav_lat']
nc_tools.show_variables(tracers)

cyc = NC.Dataset('/ocean/imachuca/MEOPAR/Ariane/results/riverparticles/times/tidal_cycle/ariane_trajectories_qualitative.nc','r');
lon_cyc=cyc.variables['traj_lon']
lat_cyc=cyc.variables['traj_lat']
dep_cyc=cyc.variables['traj_depth']
x_cyc=cyc.variables['init_x']
y_cyc=cyc.variables['init_y']
lon_cyc.shape

orig = NC.Dataset('/ocean/imachuca/MEOPAR/Ariane/results/riverparticles/compare_riverdepth/original/ariane_trajectories_qualitative.nc','r');
lon_orig=orig.variables['traj_lon']
lat_orig=orig.variables['traj_lat']
dep_orig=orig.variables['traj_depth']
x_orig=orig.variables['init_x']
y_orig=orig.variables['init_y']
lon_orig.shape

m2 = NC.Dataset('/ocean/imachuca/MEOPAR/Ariane/results/riverparticles/compare_riverdepth/2m/ariane_trajectories_qualitative.nc','r');
lon_2m=m2.variables['traj_lon']
lat_2m=m2.variables['traj_lat']
dep_2m=m2.variables['traj_depth']
x_2m=m2.variables['init_x']
y_2m=m2.variables['init_y']
lon_2m.shape

tracers1 = NC.Dataset('/ocean/nsoontie/MEOPAR/SalishSea/results/storm-surges/final/dec2006/all_forcing/restart/SalishSea_1h_20061214_20061215_grid_T.nc')
salo = tracers1.variables['vosaline']
ssho = tracers1.variables['sossheig']
timestepso = tracers1.variables['time_counter']
s_deptho = tracers1.variables['deptht']
s_lono = tracers1.variables['nav_lon']
s_lato = tracers1.variables['nav_lat']

tracers2 = NC.Dataset('/ocean/nsoontie/MEOPAR/SalishSea/results/storm-surges/final/dec2006/rivers2m/1hour/SalishSea_1h_20061214_20061215_grid_T.nc')
sal2 = tracers2.variables['vosaline']
ssh2 = tracers2.variables['sossheig']
timesteps2 = tracers2.variables['time_counter']
s_depth2 = tracers2.variables['deptht']
s_lon2 = tracers2.variables['nav_lon']
s_lat2 = tracers2.variables['nav_lat']

# Slice (in coordinates)
x_slice = np.linspace(-123.2,-123.12,num=5)
ylocn = 49.101
y_slice = ylocn*np.ones_like(x_slice)

shape1 = x_slice.shape
shape2 = y_slice.shape

shape1, x_slice, shape2, y_slice

# Conversion between indices and coordinates
ilonl = -123.2
ilonr = -123.12
ilat = ylocn
iyl, ixl = tidetools.find_closest_model_point(ilonl, ilat, X, Y, bathy, allow_land=False)
iyr, ixr = tidetools.find_closest_model_point(ilonr, ilat, X, Y, bathy, allow_land=False)
ixl,iyl,ixr,iyr


# Slice (in indices)
n=8
x_slicei = np.linspace(ixl,ixr,num=n)
y_slicei = np.linspace(iyl,iyr,num=n)

shape4 = x_slicei.shape
shape5 = y_slicei.shape
shape4, x_slicei, shape5, y_slicei

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

# Figure 1
N=5
ax1.scatter(rd_lont[0,N],rd_latt[0,N],color='0.30',marker='s')
ax1.scatter(rd_lont[1:,N],rd_latt[1:,N],color='r')
viz_tools.plot_land_mask(ax1,grid,coords='map', color=land_colour)
ax1.plot(x_slice, y_slice, linestyle='solid', linewidth=3, color='black')
ax1.set_xlim([-123.25,-123])
ax1.set_ylim([49,49.2])
ax1.set_title('Ideal Slice')
ax1.set_xlabel('lon')
ax1.set_ylabel('lat')

# Figure 2
viz_tools.plot_land_mask(ax2,grid,color=land_colour)
ax2.plot(np.round(x_slicei), np.round(y_slicei), 'o',color='black')
ax2.set_xlim([300, 340]); ax2.set_ylim([390,440])
ax2.set_title('Slice for Salinity Profile')
ax2.set_xlabel('x index')
ax2.set_ylabel('y index')

# Salinity Slice
salx= np.array((np.round(x_slicei)),int)
saly = np.array((np.round(y_slicei)),int)

fig, (ax3, ax4) = plt.subplots(1, 2, figsize=(16, 8))

# Figure 3 
land_colour = 'burlywood'
viz_tools.plot_land_mask(ax3,grid,coords='map',color=land_colour)
ax3.scatter(rd_lont[0,N],rd_latt[0,N],color='0.30',marker='s')
ax3.scatter(rd_lont[1:,N],rd_latt[1:,N],color='red')
ax3.plot(s_lon[saly,salx],s_lat[saly,salx],'ok')
ax3.set_xlim([-123.25,-123])
ax3.set_ylim([49,49.2])
ax3.set_title('Slice for Salinity Profile and Particle Trajectory')
ax3.set_xlabel('lon')
ax3.set_ylabel('lat')

# Figure 4
t=47
sali=sal[t,:,saly,salx]
mesh=ax4.contourf(sali,10)
cbar=fig.colorbar(mesh)
ax4.set_ylim([4,0])
ax4.set_xlabel('Index black dot from left to right')
ax4.set_ylabel('Vertical grid cell')
ax4.set_title('Salinity Profile')
ax4.set_xticks(np.arange(n))
ax4.set_yticks([4,3,2,1,0])

# Animation
N=5

#Empty figures
fig, (ax3, ax4) = plt.subplots(1, 2, figsize=(16, 8))

pnt = [0,1,2,3,4,5,6,7]
ii=np.arange(n)
#Initial image
def init():
    land_colour = 'burlywood'
    viz_tools.plot_land_mask(ax3,grid,coords='map',color=land_colour)
    ax3.set_xlim([-123.25,-123])
    ax3.set_ylim([49,49.2])
    ax3.plot(s_lon[saly,salx],s_lat[saly,salx],'sk',label='Slice for Salinity Profile')
    for lab, i in zip (pnt,ii):
        ax3.annotate(lab,(s_lon[saly[i],salx[i]],s_lat[saly[i],salx[i]]),fontsize=11)
    ax3.scatter(rd_lont[0,N],rd_latt[0,N],color='0.30',marker='s')
    ax3.set_title('Why is this particle trapped?')
    ax3.set_xlabel('lon')
    ax3.set_ylabel('lat')
    
t = np.arange(48)
v = np.linspace(0, 20, 15, endpoint=True)

def animate(p):
    ax3.scatter(rd_lont[p+1,N],rd_latt[p+1,N],color='red',label='Particle Trajectory')
    sali=sal[t[p],:,saly,salx]
    mesh=ax4.contourf(sali,v, cmap=plt.cm.jet)
    if p == 0:
        cbar=fig.colorbar(mesh, ticks=[0,2,4,6,8,10,12,14,16,18,20])
        ax3.legend(scatterpoints = 1)
    ax4.set_ylim([4,0])
    ax4.set_xlabel('Black Squares')
    ax4.set_ylabel('Vertical grid cell')
    ax4.set_title('Salinity Profile at Slice')
    ax4.set_xticks(np.arange(n))
    ax4.set_yticks([4,3,2,1,0])
    
 
    
#The animation function (max frames=47)   
anim = animation.FuncAnimation(fig, animate, init_func=init, frames=1, interval=300, blit=False, repeat=False)

#A line that makes it all work
mywriter = animation.FFMpegWriter()

#Save in current folder
anim.save('SaltWedge000.mp4',writer=mywriter,fps=1)

wlev = stormtools.load_observations('13-Dec-2006', '16-Dec-2006','PointAtkinson')
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 Point Atkinson')

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]

wlev = stormtools.load_observations('13-Dec-2006', '16-Dec-2006','PointAtkinson')


(fig,axs)=plt.subplots(8,2,figsize=(20,50))
tt = [17,19,22,25,30,33,36,39]
colorss = ['Maroon','MediumBlue','Olive','DarkViolet','Tomato','DeepSkyBlue','SpringGreen','DeepPink']
n=np.arange(8)

for ax,t,colors in zip(axs[:,1],tt,colorss):
    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 Point Atkinson')
    ax.plot(wlev.time[t],wlev.slev[t],'o',color=colors,markersize=9)

for ax,N,colors in zip(axs[:,0],n,colorss):
    ax.scatter(lon_cyc[0,N],lat_cyc[0,N],color='0.30',marker='s')
    ax.scatter(lon_cyc[1:,N],lat_cyc[1:,N],color=colors,marker='o')
    viz_tools.plot_land_mask(ax,grid,coords='map', color='burlywood')
    ax.set_xlim([-123.4,-123])
    ax.set_ylim([48.9,49.3])
    ax.set_xlabel('lon')
    ax.set_ylabel('lat')
    ax.set_title('Trajectory')

(fig,axs)=plt.subplots(1,3,figsize=(25,10))
n=np.arange(3)

for ax,N in zip(axs,n):
    ax.scatter(lon_orig[0,N],lat_orig[0,N],color='0.30',marker='s')
    ax.scatter(lon_orig[1:,N],lat_orig[1:,N],color='MediumBlue',marker='.')
    ax.scatter(lon_2m[1:,N],lat_2m[1:,N],color='Magenta',marker='.')
    viz_tools.plot_land_mask(ax,grid,coords='map', color='burlywood')
    ax.set_xlim([-123.4,-123])
    ax.set_ylim([48.9,49.3])
    ax.set_xlabel('lon')
    ax.set_ylabel('lat')
    ax.set_title('Trajectory')

fig, (ax3, ax4) = plt.subplots(1, 2, figsize=(16, 8))

def init():
        return [fig, ax3, ax4]

def animate(t):
    sali2=sal2[t,:,saly,salx]
    mesh2=ax4.contourf(sali2,10)
    ax4.set_ylim([4,0])
    ax4.set_xlabel('Slice')
    ax4.set_ylabel('Vertical grid cell')
    ax4.set_title('2 metres')
    ax4.set_xticks(np.arange(n))
    ax4.set_yticks([4,3,2,1,0])

    salio=salo[t,:,saly,salx]
    mesho=ax3.contourf(salio,10)
    ax3.set_ylim([4,0])
    ax3.set_xlabel('Slice')
    ax3.set_ylabel('Vertical grid cell')
    ax3.set_title('original')
    ax3.set_xticks(np.arange(n))
    ax3.set_yticks([4,3,2,1,0])

#The animation function (max frames=47)   
anim = animation.FuncAnimation(fig, animate, init_func=init, frames=48, interval=300, blit=False, repeat=False)

#A line that makes it all work
mywriter = animation.FFMpegWriter()

#Save in current folder
anim.save('SaltWedge_Compare.mp4',writer=mywriter,fps=1)