import netCDF4 as nc
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import animation
from salishsea_tools import (nc_tools)
%matplotlib inline

data = nc.Dataset('/ocean/imachuca/MEOPAR/Ariane/results/soft_carpentry/ariane_trajectories_qualitative.nc','r');

nc_tools.show_variables(data)

lon=data.variables['traj_lon']
lat=data.variables['traj_lat']
dep=data.variables['traj_depth']
x=data.variables['init_x']
y=data.variables['init_y']

lon.shape

fig, ax = plt.subplots(1,1,figsize=(20,10))
ax.scatter(lon[:,:],lat[:,:],color='blue')

ax.set_title('2D')
ax.set_xlabel('Longitude')
ax.set_ylabel('Latitude')

ax.set_xlim([-124.6,-123.0])        
ax.set_ylim([48.25,48.65])

fig = plt.figure(figsize=(20,10))
ax = fig.gca(projection='3d')
ax.scatter(lon[:,:],lat[:,:],dep[:,0],color='blue')

ax.set_title('3D')
ax.set_xlabel('Longitude')
ax.set_ylabel('Latitude')
ax.set_zlabel('Depth')

ax.set_xlim([-124.6,-123.0])        
ax.set_ylim([48.25,48.65])
ax.set_zlim([-100,0])

fig = plt.figure(figsize=(20,10))
ax = fig.gca(projection='3d')
ax.scatter(lon[:,:],lat[:,:],dep[:,:],color='blue')

ax.set_title('3D')
ax.set_xlabel('Longitude')
ax.set_ylabel('Latitude')
ax.set_zlabel('Depth')

ax.set_xlim([-124.6,-123.0])        
ax.set_ylim([48.25,48.65])
ax.set_zlim([-100,0])

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

#Empty map
fig, ax = plt.subplots(1,1,figsize=(20,10))

ax.set_title('2D')
ax.set_xlabel('Longitude')
ax.set_ylabel('Latitude')

ax.set_xlim([-124.6,-123.0])        
ax.set_ylim([48.25,48.65])

#Initial image
def init():
    return [fig, ax]

#Animated points
def animate(p):
    ax.scatter(lon[p,:],lat[p,:],color='blue')

#The animation function    
anim = animation.FuncAnimation(fig, animate, init_func=init, frames=10)

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

#Save in current folder
anim.save('Software-Carpentry2D.mp4',writer=mywriter)

#Show as a pop-up window
#plt.show()

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

ax.set_title('3D')
ax.set_xlabel('Longitude')
ax.set_ylabel('Latitude')
ax.set_zlabel('Depth')

ax.set_xlim([-124.6,-123.0])        
ax.set_ylim([48.25,48.65])
ax.set_zlim([-100,0])

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

#Initial image
def init():
    return [fig, ax]

#Animated points
def animate(p):
    ax.scatter(lon[p,:],lat[p,:],dep[p,:],color='blue')

#The animation function    
anim = animation.FuncAnimation(fig, animate, init_func=init, frames=10)

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

#Save in current folder
anim.save('Software-Carpentry3D.mp4',writer=mywriter)

#Show as a pop-up window
#plt.show()