#!/usr/bin/env python
# coding: utf-8

# # Drifter Simulations
# 
# ***
# Based on a notebook by Ben Moore-Maley and modified by Susan Allen

# In[1]:


import numpy as np
import xarray as xr
import os
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
import matplotlib.animation as animation
from datetime import datetime, timedelta
from dateutil.parser import parse
from scipy.io import loadmat
from tqdm import tqdm

from parcels import FieldSet, ParticleSet, JITParticle, ErrorCode, AdvectionRK4

get_ipython().run_line_magic('matplotlib', 'inline')


# In[2]:


plt.rcParams['font.size'] = 14


# ***
# 
# ### Local functions

# In[3]:


def make_prefix(date, path, res='h'):
    """Construct path prefix for local SalishSeaCast results given date object and paths dict
    e.g., ./SalishSea/SalishSea_1h_yyyymmdd_yyyymmdd
    """

    datestr = '_'.join(np.repeat(date.strftime('%Y%m%d'), 2))
    prefix = os.path.join(path, f'SalishSea_1{res}_{datestr}')
    
    return prefix


def DeleteParticle(particle, fieldset, time):
    """Delete particle from OceanParcels simulation to avoid run failure
    """
    
    print(f'Particle {particle.id} lost !! [{particle.lon}, {particle.lat}, {particle.depth}, {particle.time}]')
    particle.delete()


# ***
# 
# ### Load drifters and definitions

# In[4]:


# Define paths
paths = {
    'NEMO': './SalishSea/',
    'coords': './grid/coordinates_seagrid_SalishSea201702.nc',
    'mask': './grid/mesh_mask201702.nc',
    'out': './results',
}


# In[5]:


# Duration and timestep [s]
# Duration
duration = timedelta(days=5) # was initially 5 - this is the total duration, or runtime, of the particle simulation
# Timestep
dt = 90 # this is the internal numerical timestep of the simulation in seconds - apparently this should be small, and
# then the output timestep is large, so your oputfile isnt too large i beleive - from OceanParcels tutorial


# In[10]:


n = 2   # number of particles, originally 1000
# Define Gaussian point cloud in the horizontal
r = 100   # radius of particle cloud [m] - initially 10,000 (10km)
deg2m = 111000 * np.cos(50 * np.pi / 180) # what is this?
var = (r / (deg2m * 3))**2
x_offset, y_offset = np.random.multivariate_normal([0, 0], [[var, 0], [0, var]], n).T
# Set a uniform distribution in depth, from dmin to dmax
dmin = 1.
dmax = 1.
zvals = dmin + np.random.random_sample(n)*(dmax-dmin)


# In[ ]:


## For my pteropod simulation -don't run the above cell - modify it to launch my pteropod particles at assigned depths
#n = 10 # number of particles


# ***
# 
# ### Simulations

# In[7]:


# Need to run this cell twice for some reason
start = datetime(2020, 3, 15) # Day that particles initialize
# A combination of a date and a time. Attributes: year, month, day, hour, minute, second, microsecond, and tzinfo.^
# Build filenames
Ulist, Vlist, Wlist = [], [], []
for day in range(duration.days + 1): # was initially (duration.days + 3) - why?? CHANGED
    path_NEMO = make_prefix(start + timedelta(days=day), paths['NEMO'])
    print (path_NEMO)
    Ulist.append(path_NEMO + '_grid_U.nc')
    Vlist.append(path_NEMO + '_grid_V.nc')
    Wlist.append(path_NEMO + '_grid_W.nc')

# Load NEMO forcing : note, depth aware but no vertical advection, particles stay at their original depth
filenames = {
    'U': {'lon': paths['coords'], 'lat': paths['coords'], 'depth': Wlist[0], 'data': Ulist},
    'V': {'lon': paths['coords'], 'lat': paths['coords'], 'depth': Wlist[0], 'data': Vlist},
}
variables = {'U': 'vozocrtx', 'V': 'vomecrty'}
dimensions = {'lon': 'glamf', 'lat': 'gphif', 'depth': 'depthw','time': 'time_counter'}

field_set = FieldSet.from_nemo(filenames, variables, dimensions)


# In[11]:


# Set output file name.  Maybe change for each run
fn = f'matt_drifters_2particles1mdepth' + '_'.join(d.strftime('%Y%m%d') for d in [start, start+duration]) + '.nc'
outfile = os.path.join(paths['out'], fn)
print(outfile)


# In[12]:


# Execute run
clon, clat = -123.4, 49.18  # choose horizontal centre of the particle cloud # inital = -123.4, 49.18
lon, lat, z = clon + x_offset, clat + y_offset, zvals
pset = ParticleSet.from_list(field_set, JITParticle, lon=lon, lat=lat, depth=z, time=start+timedelta(hours=1))
# I think this is where I can change the hour that the particles initialize^ Can't be 0 because that puts it outside
# the timerange - I guess technically that would be the previous day at midnight?
pset.execute(
    pset.Kernel(AdvectionRK4), runtime=duration, dt=dt, # I believe this is where you'd put a - for reverse i.e. -dt
    output_file=pset.ParticleFile(name=outfile, outputdt=timedelta(hours=1)), # This is the output timestep - 1hr
    recovery={ErrorCode.ErrorOutOfBounds: DeleteParticle},
)


# In[13]:


# this cell will fail, but I seem to need to run it to get the outputfiles from the temp directory into my 
# final outfile
pset.execute(fail=fail)


# In[14]:


## Can print the particles to see thier coordinates and depth - dont do this if the number of particles is large
print(pset)


# In[25]:


## Easier way to plot? test - the U might not work because its 3D here, not just surface - not working
# pset.show(field=field_set.U, domain={'N':49.8, 'S':48.6, 'E':-122.8 ,'W':-124})


# In[15]:


#coords.close()
#mask.close()
coords = xr.open_dataset(paths['coords'], decode_times=False)
mask = xr.open_dataset(paths['mask'])


# In[16]:


ds = xr.open_dataset(outfile)


# In[17]:


print(ds)


# In[18]:


# Plot the paths of all particles
# Initialize a figure and specify the size
fig, ax = plt.subplots(figsize=(19, 8))
# Set the fill colour for the map polygons
ax.contourf(coords.nav_lon, coords.nav_lat, mask.tmask[0, 0, ...], levels=[-0.01, 0.01], colors='lightgray')
# Set the line colour for the map polygons ('k' = black)
ax.contour(coords.nav_lon, coords.nav_lat, mask.tmask[0, 0, ...], levels=[-0.01, 0.01], colors='k')
# Set axes limits in map coordinates
ax.set_xlim([-124., -122.7])
ax.set_ylim([48.55, 49.8])
# Set the aspect ratio
ax.set_aspect(5/4.4)
# Not sure what this does but doesnt seem to affect the plot with it removed (and removed from next line too)
nmin, nmax = 0, -1
# Loops through the data set in range(n) which means the size of the dataset, and s is defined as the image I think
for traj in range(n):
    s = ax.scatter(ds.lon[traj, nmin:nmax], ds.lat[traj, nmin:nmax], s = 20)

# Also need to figure out how to set the color map for only one timestep

# Can just click and drag the figure into a directory to save as a .png - how do you specify the resolution tho?
# save figure this way? fig.savefig('yourfilename.png')
# Can make multiple sub-plots within one figure onject with: fig = plt.subplots(nrows=2, ncols=2)


# In[32]:


## Plot start locations only
fig, ax = plt.subplots(figsize=(19, 8))
ax.contourf(coords.nav_lon, coords.nav_lat, mask.tmask[0, 0, ...], levels=[-0.01, 0.01], colors='lightgray')
ax.contour(coords.nav_lon, coords.nav_lat, mask.tmask[0, 0, ...], levels=[-0.01, 0.01], colors='k')
ax.set_xlim([-124., -122.7])
ax.set_ylim([48.55, 49.8])
ax.set_aspect(5/4.4)
nmin, nmax = 0, -1
# Loops through the data set in range(n) which means the size of the dataset, and s is defined as the image I think
for traj in range(n):
    s = ax.scatter(ds.lon[traj, nmin], ds.lat[traj, nmin], s = 20)


# In[33]:


## Plot end locations only
fig, ax = plt.subplots(figsize=(19, 8))
ax.contourf(coords.nav_lon, coords.nav_lat, mask.tmask[0, 0, ...], levels=[-0.01, 0.01], colors='lightgray')
ax.contour(coords.nav_lon, coords.nav_lat, mask.tmask[0, 0, ...], levels=[-0.01, 0.01], colors='k')
ax.set_xlim([-124., -122.7])
ax.set_ylim([48.55, 49.8])
ax.set_aspect(5/4.4)
nmin, nmax = 0, -1
# Loops through the data set in range(n) which means the size of the dataset, and s is defined as the image I think
for traj in range(n):
    s = ax.scatter(ds.lon[traj, nmax], ds.lat[traj, nmax], s = 20)


# In[18]:


# Make colormapped by depth - takes long time to run
# Initialize a figure and specify the size
fig, ax = plt.subplots(figsize=(19, 8))
# Set the fill colour for the map polygons
ax.contourf(coords.nav_lon, coords.nav_lat, mask.tmask[0, 0, ...], levels=[-0.01, 0.01], colors='lightgray')
# Set the line colour for the map polygons ('k' = black)
ax.contour(coords.nav_lon, coords.nav_lat, mask.tmask[0, 0, ...], levels=[-0.01, 0.01], colors='k')
# Set axes limits in map coordinates
ax.set_xlim([-124., -122.7])
ax.set_ylim([48.55, 49.8])
# Set the aspect ratio
ax.set_aspect(5/4.4)
# Why is nmax -1???
nmin, nmax = 0, -1
# Loops through the data set in range(n) which means the size of the dataset, and s is defined as the image I think
for traj in tqdm(range(n)):
    s = ax.scatter(ds.lon, ds.lat, s = 20, c = ds['z'], cmap = 'jet')
# Make colorbar for depth
cbar = plt.colorbar(s)
cbar.set_label('Depth')
cbar.ax.invert_yaxis()

# Can just click and drag the figure into a directory to save as a .png - how do you specify the resolution tho?
# save figure this way? fig.savefig('yourfilename.png')
# How to make the points smaller so I can see the paths better?
# Can make multiple sub-plots within one figure onject with: fig = plt.subplots(nrows=2, ncols=2)


# In[30]:


nmax = -1
fig, ax = plt.subplots(figsize=(10, 10))
for traj in tqdm(range(n)):
    ax.plot(ds.lon[traj, nmax], -ds.z[traj, nmax], 'o')


# In[20]:


# Make same fig as above but with scatter instead of plot
fig, ax = plt.subplots(figsize=(10, 10))
for traj in range(n):
    ax.scatter(ds.lon[traj, nmax], -ds.z[traj, nmax], s = None)


# In[22]:


# Trying to make vertical slice colormapped by depth - not working
nmax = -1
fig, ax = plt.subplots(figsize=(10, 10))
for traj in tqdm(range(n)):
    ax.scatter(ds['lon'], -ds['z'], s = None, c = ds['z'], cmap = 'jet')
# ax.scatter(ds['lon'], -ds['z'], s = None, c = ds['z'], cmap = 'jet') works the same as:
# ax.scatter(ds.lon, -ds.z, s = None, c = ds['z'], cmap = 'jet')
ax.grid()
cbar = plt.colorbar(s)
cbar.set_label('Depth')
cbar.ax.invert_yaxis()


# In[287]:


# NEED to run this to close the data set before tweaking it for another run - do I tho? I think just renaming the 
# outfile fixes that
ds.close()


# In[25]:


# Attempt to animate

# Time index
timerange = ['2020 Mar 15 02:00', '2020 Mar 22 02:00']
time_ind = parse(timerange[0])

# nmin index
nmin, nmax = 0 , -1

fig, ax = plt.subplots(figsize=(19, 8))
ax.contourf(coords.nav_lon, coords.nav_lat, mask.tmask[0, 0, ...], levels=[-0.01, 0.01], colors='lightgray')
ax.contour(coords.nav_lon, coords.nav_lat, mask.tmask[0, 0, ...], levels=[-0.01, 0.01], colors='k')
ax.set_xlim([-124., -122.7])
ax.set_ylim([48.55, 49.8])
ax.set_aspect(5/4.4)

for traj in range(n):
    s = ax.scatter(ds.lon[traj, nmin], ds.lat[traj, nmin], s = 20, c = ds['z'], cmap = 'jet') # remove colormapping?

# Add timestamp
TEXT_OBJ = ax.text(0.02, 1.02, time_ind.strftime('%a %Y-%m-%d %H:%M:%S'), transform=ax.transAxes) 
    
# Also need to figure out how to set the color map for only one timestep

# Make colorbar for depth
#cbar = plt.colorbar(s)
#cbar.set_label('Depth')
#cbar.ax.invert_yaxis()



# In[26]:


def make_figure(time_ind, nmin, nmax, ax, s, ds):
    """
    """
    
    for traj in range(n):
        s = ax.scatter(ds.lon[traj, nmin], ds.lat[traj, nmin], s = 20)
    
    # Add timestamp
    TEXT_OBJ = ax.text(0.02, 1.02, time_ind.strftime('%a %Y-%m-%d %H:%M:%S'), transform=ax.transAxes)
    
    # Return plot objects as dictionary
    PLOT_OBJS = {'s': s, 'TEXT_OBJ': TEXT_OBJ}
    
    return ds, PLOT_OBJS


# In[27]:


# Time index
time_ind = parse(timerange[0])
# nmin index
nmin, nmax = 0, -1

# Make figure window
fig, ax = plt.subplots(figsize=(19, 8))
ax.contourf(coords.nav_lon, coords.nav_lat, mask.tmask[0, 0, ...], levels=[-0.01, 0.01], colors='lightgray')
ax.contour(coords.nav_lon, coords.nav_lat, mask.tmask[0, 0, ...], levels=[-0.01, 0.01], colors='k')
ax.set_xlim([-124., -122.7])
ax.set_ylim([48.55, 49.8])
ax.set_aspect(5/4.4)

# Make figure
ds, PLOT_OBJS = make_figure(time_ind, nmin, nmax, ax, s, ds)

plt.show()


# In[28]:


# Time index
starttime = parse(timerange[0])
# nmin index
newnmin, nmax = 0, -1

# Make figure window
fig, ax = plt.subplots(figsize=(19, 8))
ax.contourf(coords.nav_lon, coords.nav_lat, mask.tmask[0, 0, ...], levels=[-0.01, 0.01], colors='lightgray')
ax.contour(coords.nav_lon, coords.nav_lat, mask.tmask[0, 0, ...], levels=[-0.01, 0.01], colors='k')
ax.set_xlim([-124., -122.7])
ax.set_ylim([48.55, 49.8])
ax.set_aspect(5/4.4)

# Make figure
ds, PLOT_OBJS = make_figure(starttime, newnmin, nmax, ax, s, ds)

# Step time index forward
time_ind = starttime + timedelta(hours=1)
# Step nmin index forward
nmin = newnmin + 1

# Clear previous particle locations - I'm not using contours so this doesnt work...
#for C in PLOT_OBJS['s'].collections:
    #C.remove()

# Update particle locations
for traj in range(n):
        PLOT_OBJS['s'] = ax.scatter(ds.lon[traj, nmin],
                       ds.lat[traj, nmin], s = 20)

# Update timestamp
PLOT_OBJS['TEXT_OBJ'].set_text(time_ind.strftime('%a %Y-%m-%d %H:%M:%S'))

plt.show()


# In[29]:


def update_figure(time_ind, nmin, nmax, ax, s, ds, PLOT_OBJS):
    """
    """
    
    # Update particle locations
    for traj in range(n):
        PLOT_OBJS['s'] = ax.scatter(ds.lon[traj, nmin],
                       ds.lat[traj, nmin], s = 20)
    
    # Update timestamp
    PLOT_OBJS['TEXT_OBJ'].set_text(time_ind.strftime('%a %Y-%m-%d %H:%M:%S'))
    
    return PLOT_OBJS


# In[30]:


# Time index
starttime = parse(timerange[0])
# nmin index
newnmin, nmax = 0, -1

# Make figure window
fig, ax = plt.subplots(figsize=(19, 8))
ax.contourf(coords.nav_lon, coords.nav_lat, mask.tmask[0, 0, ...], levels=[-0.01, 0.01], colors='lightgray')
ax.contour(coords.nav_lon, coords.nav_lat, mask.tmask[0, 0, ...], levels=[-0.01, 0.01], colors='k')
ax.set_xlim([-124., -122.7])
ax.set_ylim([48.55, 49.8])
ax.set_aspect(5/4.4)

# Make figure
ds, PLOT_OBJS = make_figure(starttime, newnmin, nmax, ax, s, ds)

# Step time index forward
time_ind = starttime + timedelta(hours=1)
# Step nmin index forward
nmin = newnmin + 1

# Update figure
PLOT_OBJS = update_figure(time_ind, nmin, nmax, ax, s, ds, PLOT_OBJS)

plt.show()


# In[31]:


# Time index
starttime = parse(timerange[0])
# nmin index
newnmin, nmax = 0, -1

# Make figure window
fig, ax = plt.subplots(figsize=(19, 8))
ax.contourf(coords.nav_lon, coords.nav_lat, mask.tmask[0, 0, ...], levels=[-0.01, 0.01], colors='lightgray')
ax.contour(coords.nav_lon, coords.nav_lat, mask.tmask[0, 0, ...], levels=[-0.01, 0.01], colors='k')
ax.set_xlim([-124., -122.7])
ax.set_ylim([48.55, 49.8])
ax.set_aspect(5/4.4)

# Make figure
ds, PLOT_OBJS = make_figure(starttime, newnmin, nmax, ax, s, ds)

def next_frame(t, PLOT_OBJS):
    # Step time index forward
    time_ind = starttime + timedelta(hours=t)
    # Step nmin index forward
    nmin = newnmin + t

    # Update figure
    PLOT_OBJS = update_figure(time_ind, nmin, nmax, ax, s, ds, PLOT_OBJS)
    
    return PLOT_OBJS

# Animate
ANI = animation.FuncAnimation(fig, next_frame, fargs=[PLOT_OBJS], frames=169) # update the number of frames?
mywriter = animation.FFMpegWriter(fps=12, bitrate=10000)
ANI.save('./Plots/mymovie8.mp4', writer=mywriter)


# Now try to place the particles in the JdF and run the simulation backwards.
# First try to only simulate one timestep or something, so that I can quickly check the starting locations and depths before running the full simulation

# In[8]:


# Duration and timestep [s]
# Duration
duration = timedelta(days = 7) # was initially days=7, this is the total duration, or runtime, of the particle simulation
# Timestep
dt = 90 # this is the internal numerical timestep of the simulation in seconds - apparently this should be small, and
# then the output timestep is large, so your oputfile isnt too large i beleive - from OceanParcels tutorial


# In[9]:


n = 1000   # number of particles, originally 1000
# Define Gaussian point cloud in the horizontal
r = 8000   # radius of particle cloud [m] - initially 10,000 (10km)
deg2m = 111000 * np.cos(50 * np.pi / 180) # what is this?
var = (r / (deg2m * 3))**2
x_offset, y_offset = np.random.multivariate_normal([0, 0], [[var, 0], [0, var]], n).T
# Set a uniform distribution in depth, from dmin to dmax
dmin = 0.
dmax = 200.
zvals = dmin + np.random.random_sample(n)*(dmax-dmin)


# In[10]:


# Need to run this cell twice for some reason
start = datetime(2020, 3, 15) # Day that particles initialize
# A combination of a date and a time. Attributes: year, month, day, hour, minute, second, microsecond, and tzinfo.^
# Build filenames
Ulist, Vlist, Wlist = [], [], []
for day in range(duration.days + 1): # was initially (duration.days + 3) - why?? CHANGED for some reason you can't 
    # make the runtime the full 8 days, you have to do 7 and add a day here...
    path_NEMO = make_prefix(start + timedelta(days=day), paths['NEMO'])
    print (path_NEMO)
    Ulist.append(path_NEMO + '_grid_U.nc')
    Vlist.append(path_NEMO + '_grid_V.nc')
    Wlist.append(path_NEMO + '_grid_W.nc')

# Load NEMO forcing : note, depth aware but no vertical advection, particles stay at their original depth
filenames = {
    'U': {'lon': paths['coords'], 'lat': paths['coords'], 'depth': Wlist[0], 'data': Ulist},
    'V': {'lon': paths['coords'], 'lat': paths['coords'], 'depth': Wlist[0], 'data': Vlist},
}
variables = {'U': 'vozocrtx', 'V': 'vomecrty'}
dimensions = {'lon': 'glamf', 'lat': 'gphif', 'depth': 'depthw','time': 'time_counter'}

field_set = FieldSet.from_nemo(filenames, variables, dimensions)


# In[11]:


# Set output file name.  Maybe change for each run
fn = f'matt_drifters_JdF1000particles0to200m' + '_'.join(d.strftime('%Y%m%d') for d in [start, start+duration]) + '.nc'
outfile = os.path.join(paths['out'], fn)
print(outfile)


# In[203]:


# Execute run
clon, clat = -123.72, 48.24  # choose horizontal centre of the particle cloud # inital = -123.4, 49.18
lon, lat, z = clon + x_offset, clat + y_offset, zvals
pset = ParticleSet.from_list(field_set, JITParticle, lon=lon, lat=lat, depth=z) # time=start+timedelta(hours=1)
# I think this is where I can change the hour that the particles initialize^ Can't be 0 because that puts it outside
# the timerange - I guess technically that would be the previous day at midnight?
pset.execute(
    pset.Kernel(AdvectionRK4), runtime=duration, dt=-dt, # I believe this is where you'd put a - for reverse i.e. -dt
    output_file=pset.ParticleFile(name=outfile, outputdt=timedelta(hours=1)), # This is the output timestep - 1hr
    recovery={ErrorCode.ErrorOutOfBounds: DeleteParticle},
)


# In[204]:


# this cell will fail, but I seem to need to run it to get the outputfiles from the temp directory into my 
# final outfile
pset.execute(fail=fail)


# In[15]:


#coords.close()
#mask.close()
coords = xr.open_dataset(paths['coords'], decode_times=False)
mask = xr.open_dataset(paths['mask'])


# In[16]:


ds = xr.open_dataset(outfile)


# In[17]:


print(ds)


# In[18]:


# NEED to run this to close the data set before tweaking it for another run, or I can just rename the outfile
ds.close()


# In[19]:


## Plot start locations only
fig, ax = plt.subplots(figsize=(19, 8))
ax.contourf(coords.nav_lon, coords.nav_lat, mask.tmask[0, 0, ...], levels=[-0.01, 0.01], colors='lightgray')
ax.contour(coords.nav_lon, coords.nav_lat, mask.tmask[0, 0, ...], levels=[-0.01, 0.01], colors='k')
ax.set_xlim([-124., -122.15])
ax.set_ylim([47.8, 49.2])
ax.set_aspect(5/4.4)
nmin, nmax = 0, -1
# Loops through the data set in range(n) which means the size of the dataset, and s is defined as the image I think
for traj in range(n):
    s = ax.scatter(ds.lon[traj, nmin], ds.lat[traj, nmin], s = 20)


# In[20]:


## Plot end locations only
fig, ax = plt.subplots(figsize=(19, 8))
ax.contourf(coords.nav_lon, coords.nav_lat, mask.tmask[0, 0, ...], levels=[-0.01, 0.01], colors='lightgray')
ax.contour(coords.nav_lon, coords.nav_lat, mask.tmask[0, 0, ...], levels=[-0.01, 0.01], colors='k')
ax.set_xlim([-124., -122.15])
ax.set_ylim([47.8, 49.2])
ax.set_aspect(5/4.4)
nmin, nmax = 0, -1
# Loops through the data set in range(n) which means the size of the dataset, and s is defined as the image I think
for traj in range(n):
    s = ax.scatter(ds.lon[traj, nmax], ds.lat[traj, nmax], s = 20)


# In[21]:


# Plot entire paths
## Plot end locations only
fig, ax = plt.subplots(figsize=(19, 8))
ax.contourf(coords.nav_lon, coords.nav_lat, mask.tmask[0, 0, ...], levels=[-0.01, 0.01], colors='lightgray')
ax.contour(coords.nav_lon, coords.nav_lat, mask.tmask[0, 0, ...], levels=[-0.01, 0.01], colors='k')
ax.set_xlim([-124., -122.15])
ax.set_ylim([47.8, 49.2])
ax.set_aspect(5/4.4)
nmin, nmax = 0, -1
# Loops through the data set in range(n) which means the size of the dataset, and s is defined as the image I think
for traj in range(n):
    s = ax.scatter(ds.lon[traj, nmin:nmax], ds.lat[traj, nmin:nmax], s = 20)


# In[50]:


# Plot starting depth profile of particles
nmax = -1
fig, ax = plt.subplots(figsize=(10, 10))
ax.set_xlim([-125., -122.])
for traj in tqdm(range(n)):
    ax.plot(ds.lon[traj, nmin], -ds.z[traj, nmin], 'o')


# In[49]:


# Plot end depth profile of particles
nmax = -1
fig, ax = plt.subplots(figsize=(10, 10))
ax.set_xlim([-125., -122.])
for traj in tqdm(range(n)):
    ax.plot(ds.lon[traj, nmin:nmax], -ds.z[traj, nmin:nmax], 'o')


# In[214]:


# Make colormapped by depth - takes long time to run
fig, ax = plt.subplots(figsize=(19, 8))
ax.contourf(coords.nav_lon, coords.nav_lat, mask.tmask[0, 0, ...], levels=[-0.01, 0.01], colors='lightgray')
ax.contour(coords.nav_lon, coords.nav_lat, mask.tmask[0, 0, ...], levels=[-0.01, 0.01], colors='k')
ax.set_xlim([-124., -122.15])
ax.set_ylim([47.8, 49.2])
ax.set_aspect(5/4.4)
nmin, nmax = 0, -1
for traj in tqdm(range(n)):
    s = ax.scatter(ds.lon, ds.lat, s = 20, c = ds['z'], cmap = 'jet')
# Make colorbar for depth
cbar = plt.colorbar(s)
cbar.set_label('Depth')
cbar.ax.invert_yaxis()


# In[22]:


# Attempt to animate

# Time index
timerange = ['2020 Mar 22 00:00', '2020 Mar 15 00:00']
time_ind = parse(timerange[0])

# nmin index
nmin, nmax = 0 , -1

fig, ax = plt.subplots(figsize=(19, 8))
ax.contourf(coords.nav_lon, coords.nav_lat, mask.tmask[0, 0, ...], levels=[-0.01, 0.01], colors='lightgray')
ax.contour(coords.nav_lon, coords.nav_lat, mask.tmask[0, 0, ...], levels=[-0.01, 0.01], colors='k')
ax.set_xlim([-124., -122.15])
ax.set_ylim([47.8, 49.2])
ax.set_aspect(5/4.4)

for traj in range(n):
    s = ax.scatter(ds.lon[traj, nmin], ds.lat[traj, nmin], s = 20)

# Add timestamp
TEXT_OBJ = ax.text(0.02, 1.02, time_ind.strftime('%a %Y-%m-%d %H:%M:%S'), transform=ax.transAxes) 
    
# Also need to figure out how to set the color map for only one timestep

# Make colorbar for depth
#cbar = plt.colorbar(s)
#cbar.set_label('Depth')
#cbar.ax.invert_yaxis()


# In[23]:


def make_figure(time_ind, nmin, nmax, ax, s, ds):
    """
    """
    
    for traj in range(n):
        s = ax.scatter(ds.lon[traj, nmin], ds.lat[traj, nmin], s = 20)
    
    # Add timestamp
    TEXT_OBJ = ax.text(0.02, 1.02, time_ind.strftime('%a %Y-%m-%d %H:%M:%S'), transform=ax.transAxes)
    
    # Return plot objects as dictionary
    PLOT_OBJS = {'s': s, 'TEXT_OBJ': TEXT_OBJ}
    
    return ds, PLOT_OBJS


# In[24]:


# Time index
time_ind = parse(timerange[0])
# nmin index
nmin, nmax = 0, -1

# Make figure window
fig, ax = plt.subplots(figsize=(19, 8))
ax.contourf(coords.nav_lon, coords.nav_lat, mask.tmask[0, 0, ...], levels=[-0.01, 0.01], colors='lightgray')
ax.contour(coords.nav_lon, coords.nav_lat, mask.tmask[0, 0, ...], levels=[-0.01, 0.01], colors='k')
ax.set_xlim([-124., -122.15])
ax.set_ylim([47.8, 49.2])
ax.set_aspect(5/4.4)

# Make figure
ds, PLOT_OBJS = make_figure(time_ind, nmin, nmax, ax, s, ds)

plt.show()


# In[25]:


# Time index
starttime = parse(timerange[0])
# nmin index
newnmin, nmax = 0, -1

# Make figure window
fig, ax = plt.subplots(figsize=(19, 8))
ax.contourf(coords.nav_lon, coords.nav_lat, mask.tmask[0, 0, ...], levels=[-0.01, 0.01], colors='lightgray')
ax.contour(coords.nav_lon, coords.nav_lat, mask.tmask[0, 0, ...], levels=[-0.01, 0.01], colors='k')
ax.set_xlim([-124., -122.15])
ax.set_ylim([47.8, 49.2])
ax.set_aspect(5/4.4)

# Make figure
ds, PLOT_OBJS = make_figure(starttime, newnmin, nmax, ax, s, ds)

# Step time index forward
time_ind = starttime - timedelta(hours=1)
# Step nmin index forward
nmin = newnmin + 1

# Clear previous particle locations - I'm not using contours so this doesnt work...
#for C in PLOT_OBJS['s'].collections:
    #C.remove()

# Update particle locations
for traj in range(n):
        PLOT_OBJS['s'] = ax.scatter(ds.lon[traj, nmin],
                       ds.lat[traj, nmin], s = 20)

# Update timestamp
PLOT_OBJS['TEXT_OBJ'].set_text(time_ind.strftime('%a %Y-%m-%d %H:%M:%S'))

plt.show()


# In[26]:


def update_figure(time_ind, nmin, nmax, ax, s, ds, PLOT_OBJS):
    """
    """
    
    # Update particle locations
    for traj in range(n):
        PLOT_OBJS['s'] = ax.scatter(ds.lon[traj, nmin],
                       ds.lat[traj, nmin], s = 20)
    
    # Update timestamp
    PLOT_OBJS['TEXT_OBJ'].set_text(time_ind.strftime('%a %Y-%m-%d %H:%M:%S'))
    
    return PLOT_OBJS


# In[27]:


# Time index
starttime = parse(timerange[0])
# nmin index
newnmin, nmax = 0, -1

# Make figure window
fig, ax = plt.subplots(figsize=(19, 8))
ax.contourf(coords.nav_lon, coords.nav_lat, mask.tmask[0, 0, ...], levels=[-0.01, 0.01], colors='lightgray')
ax.contour(coords.nav_lon, coords.nav_lat, mask.tmask[0, 0, ...], levels=[-0.01, 0.01], colors='k')
ax.set_xlim([-124., -122.15])
ax.set_ylim([47.8, 49.2])
ax.set_aspect(5/4.4)

# Make figure
ds, PLOT_OBJS = make_figure(starttime, newnmin, nmax, ax, s, ds)

# Step time index forward
time_ind = starttime - timedelta(hours=1)
# Step nmin index forward
nmin = newnmin + 1

# Update figure
PLOT_OBJS = update_figure(time_ind, nmin, nmax, ax, s, ds, PLOT_OBJS)

plt.show()


# In[28]:


# Time index
starttime = parse(timerange[0])
# nmin index
newnmin, nmax = 0, -1

# Make figure window
fig, ax = plt.subplots(figsize=(19, 8))
ax.contourf(coords.nav_lon, coords.nav_lat, mask.tmask[0, 0, ...], levels=[-0.01, 0.01], colors='lightgray')
ax.contour(coords.nav_lon, coords.nav_lat, mask.tmask[0, 0, ...], levels=[-0.01, 0.01], colors='k')
ax.set_xlim([-124., -122.15])
ax.set_ylim([47.8, 49.2])
ax.set_aspect(5/4.4)

# Make figure
ds, PLOT_OBJS = make_figure(starttime, newnmin, nmax, ax, s, ds)

def next_frame(t, PLOT_OBJS):
    # Step time index forward
    time_ind = starttime - timedelta(hours=t)
    # Step nmin index forward
    nmin = newnmin + t

    # Update figure
    PLOT_OBJS = update_figure(time_ind, nmin, nmax, ax, s, ds, PLOT_OBJS)
    
    return PLOT_OBJS

# Animate
ANI = animation.FuncAnimation(fig, next_frame, fargs=[PLOT_OBJS], frames=168) # update the number of frames?
mywriter = animation.FFMpegWriter(fps=12, bitrate=10000)
ANI.save('./Plots/mymovie10.mp4', writer=mywriter)


# In[84]:


# Attempt to animate the vertical slice depth profile of particles
    
# Time index
timerange = ['2020 Mar 22 00:00', '2020 Mar 15 00:00'] # I don't think this start time is correct - maybe 23:00?
time_ind = parse(timerange[0])

# nmin index
nmin, nmax = 0 , -1

fig, ax = plt.subplots(figsize=(10, 10))
ax.set_xlim([-125., -122.])
for traj in range(n):
    s = ax.plot(ds.lon[traj, nmin], -ds.z[traj, nmin], 'o')

# Add timestamp
TEXT_OBJ = ax.text(0.02, 1.02, time_ind.strftime('%a %Y-%m-%d %H:%M:%S'), transform=ax.transAxes)


# In[85]:


def make_figure(time_ind, nmin, nmax, ax, s, ds):
    """
    """
    
    for traj in range(n):
        s = ax.plot(ds.lon[traj, nmin], -ds.z[traj, nmin], 'o') # this just plots the very first frame
    
    # Add timestamp
    TEXT_OBJ = ax.text(0.02, 1.02, time_ind.strftime('%a %Y-%m-%d %H:%M:%S'), transform=ax.transAxes)
    
    # Return plot objects as dictionary
    PLOT_OBJS = {'s': s, 'TEXT_OBJ': TEXT_OBJ}
    
    return ds, PLOT_OBJS


# In[86]:


# Time index
time_ind = parse(timerange[0])
# nmin index
nmin, nmax = 0, -1

# Make figure window
fig, ax = plt.subplots(figsize=(10, 10))
ax.set_xlim([-125., -122.])

# Make figure
ds, PLOT_OBJS = make_figure(time_ind, nmin, nmax, ax, s, ds)

plt.show()


# In[97]:


# Time index
starttime = parse(timerange[0])
# nmin index
newnmin, nmax = 0, -1

# Make figure window
fig, ax = plt.subplots(figsize=(10, 10))
ax.set_xlim([-125., -122.])

# Make figure
ds, PLOT_OBJS = make_figure(starttime, newnmin, nmax, ax, s, ds)

# Step time index forward
time_ind = starttime - timedelta(hours=1)
# Step nmin index forward
nmin = newnmin + 1

# Update particle locations
for traj in range(n):
        PLOT_OBJS['s'] = ax.plot(ds.lon[traj, nmin], -ds.z[traj, nmin], 'o')
        
# Update timestamp
PLOT_OBJS['TEXT_OBJ'].set_text(time_ind.strftime('%a %Y-%m-%d %H:%M:%S'))

plt.show()


# In[98]:


def update_figure(time_ind, nmin, nmax, ax, s, ds, PLOT_OBJS):
    """
    """
        
    # Update particle locations
    for traj in range(n):
        s = ax.plot(ds.lon[traj, nmin], -ds.z[traj, nmin], 'o')
    
    # Update timestamp
    PLOT_OBJS['TEXT_OBJ'].set_text(time_ind.strftime('%a %Y-%m-%d %H:%M:%S'))
    
    return PLOT_OBJS


# In[99]:


# Time index
starttime = parse(timerange[0])
# nmin index
newnmin, nmax = 0, -1

# Make figure window
fig, ax = plt.subplots(figsize=(10, 10))
ax.set_xlim([-125., -122.])

# Make figure
ds, PLOT_OBJS = make_figure(starttime, newnmin, nmax, ax, s, ds)

# Step time index forward
time_ind = starttime - timedelta(hours=1)
# Step nmin index forward
nmin = newnmin + 1

# Update figure
PLOT_OBJS = update_figure(time_ind, nmin, nmax, ax, s, ds, PLOT_OBJS)

plt.show()


# In[100]:


# Time index
starttime = parse(timerange[0])
# nmin index
newnmin, nmax = 0, -1

# Make figure window
fig, ax = plt.subplots(figsize=(10, 10))
ax.set_xlim([-125., -122.])

# Make figure
ds, PLOT_OBJS = make_figure(starttime, newnmin, nmax, ax, s, ds)

def next_frame(t, PLOT_OBJS):
    # Step time index forward
    time_ind = starttime - timedelta(hours=t)
    # Step nmin index forward
    nmin = newnmin + t
        
    # Update figure
    PLOT_OBJS = update_figure(time_ind, nmin, nmax, ax, s, ds, PLOT_OBJS)
    
    return PLOT_OBJS

# Animate
ANI = animation.FuncAnimation(fig, next_frame, fargs=[PLOT_OBJS], frames=12) # update the number of frames?
mywriter = animation.FFMpegWriter(fps=12, bitrate=10000)
ANI.save('./Plots/mymovie14.mp4', writer=mywriter)


# In[101]:


## Attempt to make the animation without plotting the previous frame or something - might need to get rid of the dictionary

# Time index
timerange = ['2020 Mar 22 00:00', '2020 Mar 15 00:00'] # I don't think this start time is correct - maybe 23:00?
time_ind = parse(timerange[0])

# nmin index
nmin, nmax = 0 , -1

fig, ax = plt.subplots(figsize=(10, 10))
ax.set_xlim([-125., -122.])
for traj in range(n):
    s = ax.plot(ds.lon[traj, nmin], -ds.z[traj, nmin], 'o')

# Add timestamp
TEXT_OBJ = ax.text(0.02, 1.02, time_ind.strftime('%a %Y-%m-%d %H:%M:%S'), transform=ax.transAxes)


# In[102]:


def make_figure(time_ind, nmin, nmax, ax, s, ds):
    """
    """
    
    for traj in range(n):
        s = ax.plot(ds.lon[traj, nmin], -ds.z[traj, nmin], 'o') # this just plots the very first frame
    
    # Add timestamp
    TEXT_OBJ = ax.text(0.02, 1.02, time_ind.strftime('%a %Y-%m-%d %H:%M:%S'), transform=ax.transAxes)
    
    # Return plot objects as dictionary
    # PLOT_OBJS = {'s': s, 'TEXT_OBJ': TEXT_OBJ}
    
    return ds, s


# In[103]:


# Time index
time_ind = parse(timerange[0])
# nmin index
nmin, nmax = 0, -1

# Make figure window
fig, ax = plt.subplots(figsize=(10, 10))
ax.set_xlim([-125., -122.])

# Make figure
ds, s = make_figure(time_ind, nmin, nmax, ax, s, ds)

plt.show()


# In[ ]:


# Time index
starttime = parse(timerange[0])
# nmin index
newnmin, nmax = 0, -1

# Make figure window
fig, ax = plt.subplots(figsize=(10, 10))
ax.set_xlim([-125., -122.])

# Make figure
ds, s = make_figure(starttime, newnmin, nmax, ax, s, ds)

# Step time index forward
time_ind = starttime - timedelta(hours=1)
# Step nmin index forward
nmin = newnmin + 1

# Clear previous points not working
s.remove()

# Update particle locations
for traj in range(n):
        s = ax.plot(ds.lon[traj, nmin], -ds.z[traj, nmin], 'o')
        
# Update timestamp
TEXT_OBJ.set_text(time_ind.strftime('%a %Y-%m-%d %H:%M:%S'))

plt.show()


# In[ ]:


def update_figure(time_ind, nmin, nmax, ax, s, ds, PLOT_OBJS):
    """
    """
        
    # Update particle locations
    for traj in range(n):
        s = ax.plot(ds.lon[traj, nmin], -ds.z[traj, nmin], 'o')
    
    # Update timestamp
    PLOT_OBJS['TEXT_OBJ'].set_text(time_ind.strftime('%a %Y-%m-%d %H:%M:%S'))
    
    return PLOT_OBJS


# In[ ]:


# Time index
starttime = parse(timerange[0])
# nmin index
newnmin, nmax = 0, -1

# Make figure window
fig, ax = plt.subplots(figsize=(10, 10))
ax.set_xlim([-125., -122.])

# Make figure
ds, PLOT_OBJS = make_figure(starttime, newnmin, nmax, ax, s, ds)

# Step time index forward
time_ind = starttime - timedelta(hours=1)
# Step nmin index forward
nmin = newnmin + 1

# Update figure
PLOT_OBJS = update_figure(time_ind, nmin, nmax, ax, s, ds, PLOT_OBJS)

plt.show()


# In[ ]:


# Time index
starttime = parse(timerange[0])
# nmin index
newnmin, nmax = 0, -1

# Make figure window
fig, ax = plt.subplots(figsize=(10, 10))
ax.set_xlim([-125., -122.])

# Make figure
ds, PLOT_OBJS = make_figure(starttime, newnmin, nmax, ax, s, ds)

def next_frame(t, PLOT_OBJS):
    # Step time index forward
    time_ind = starttime - timedelta(hours=t)
    # Step nmin index forward
    nmin = newnmin + t
        
    # Update figure
    PLOT_OBJS = update_figure(time_ind, nmin, nmax, ax, s, ds, PLOT_OBJS)
    
    return PLOT_OBJS

# Animate
ANI = animation.FuncAnimation(fig, next_frame, fargs=[PLOT_OBJS], frames=12) # update the number of frames?
mywriter = animation.FFMpegWriter(fps=12, bitrate=10000)
ANI.save('./Plots/mymovie14.mp4', writer=mywriter)