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

# Matt's parcels tests

# In[1]:


# imports modules and renames them short names
import numpy as np
import xarray as xr
import os
from matplotlib import pyplot as plt, animation
from datetime import datetime, timedelta
from dateutil.parser import parse
from IPython.display import HTML
from salishsea_tools import nc_tools, places

# imports functions from the parcels module
from parcels import FieldSet, Field, VectorField, ParticleSet, JITParticle, ErrorCode, AdvectionRK4, AdvectionRK4_3D, plotTrajectoriesFile

# makes the plots show up below the code cells
get_ipython().run_line_magic('matplotlib', 'inline')


# In[2]:


# Sets the default font size of matplotlib plots
plt.rcParams['font.size'] = 12


# Fieldset functions

# In[3]:


def fieldset_from_nemo(daterange, coords):
    """Generate a fieldset from a hourly SalishSeaCast forcing fields
    over daterange.
    """

    # Generate sequential list of forcing file prefixes
    prefixes = [
        nc_tools.get_hindcast_prefix(daterange[0] + timedelta(days=d)) # This uses the get_hindcast_prefix function which I think is from SalishSeaTools and was made by the MOAD group. It sets the prefix with results/salishsea... etc. so you dont have to specify that
        for d in range(np.diff(daterange)[0].days + 1)
    ]

    # Predefine fieldset argument dictionaries
    filenames, variables, dimensions = {}, {}, {}

    # Define dict fields for each variable. This might be where I would add other variables from the model
    for var, name in zip(['U', 'V', 'W'], ['vozocrtx', 'vomecrty', 'vovecrtz']):

        # Dict of filenames containing the coordinate and forcing variables
        datafiles = [prefix + f'_grid_{var}.nc' for prefix in prefixes]
        filenames[var] = {'lon': coords, 'lat': coords, 'data': datafiles}

        # NEMO variable name
        variables[var] = name

        # Dict of NEMO coordinate names (f-points)
        dimensions[var] = {'lon': 'glamf', 'lat': 'gphif', 'time': 'time_counter'}
        
        # Add depth fields (f-points are on W grid)
        filenames[var]['depth'] = prefixes[0] + '_grid_W.nc'
        dimensions[var]['depth'] = 'depthw'

    # Load NEMO forcing into fieldset
    field_set = FieldSet.from_nemo(filenames, variables, dimensions, allow_time_extrapolation=True)
    
    return field_set


# Kernels

# In[4]:


# I think this causes any particle that goes outside the boundary to be deleted
def DeleteParticle(particle, fieldset, time):
    print(f'Particle {particle.id} lost !! [{particle.lon}, {particle.lat}, {particle.depth}, {particle.time}]')
    particle.delete()


# Simulations

# In[5]:


# Paths and filenames
paths = {
    'coords': '/data/SalishSeaCast/grid/coordinates_seagrid_SalishSea201702.nc',
    'mask': '/data/SalishSeaCast/grid/mesh_mask201702.nc',
    'results': '/ocean/mattmiller/MOAD/results/parcels/test',
}

# Load coords and mask files and extract grid variables
coords, mask = [xr.open_dataset(paths[key], decode_times=False) for key in ('coords', 'mask')]
gridlon, gridlat = [coords[key][0, ...].values for key in ('glamt', 'gphit')]
tmask = mask.tmask[0, 0, ...].values

# Define release parameters
location = 'Strait of Georgia' # there are predefined location names built into SalishSeaTools, under salishsea_tools.places.PLACES which can be found easily on github
n = 100   # number of particles
r = 50   # radius of particle cloud [m]

# Start time, duration and timestep
start = datetime(2019, 1, 1, 12, 30, 0) # year, month, day, hour, minute, second
duration = timedelta(days=3) # the total duration of the simulation
dt = timedelta(seconds=90) # the timestep - toggle between + and - to run forwards or reverse
# is this the timestep that parcels uses or of the fieldset from SalishSeaCast?

# Create Gaussian distribution around release point
mean, cov = [0, 0], [[r**2, 0], [0, r**2]]
x_offset, y_offset = np.random.multivariate_normal(mean, cov, n).T
lon, lat = places.PLACES[location]['lon lat']
lons = lon + x_offset / 111000 / np.cos(np.deg2rad(lat)) # this code might be used to convert from SalishSeaCast's grid to a lat lon grid, since it is an angled domain grid
lats = lat + y_offset / 111000

# Forcing daterange (I add 1-day buffers)
daterange = [start - timedelta(days=1), start + duration + timedelta(days=1)]

# Output file prefix
strings = [location] + [t.strftime('%Y%m%dT%H%M%S') for t in (start, start + duration)]
prefix = os.path.join(paths['results'], '_'.join(strings))


# Particle simulations (3D)
# 
# Build forcing fieldset

# In[6]:


# Load SalishSeaCast results into fieldset
fieldset = fieldset_from_nemo(daterange, paths['coords'])


# Run simulation with NEMO forcing - only need to do once with current particle settings, then can load the particle file in the plot codes below

# In[30]:


# Execute NEMO-only, 3D run, release at 5m
pset = ParticleSet.from_list(fieldset, JITParticle, lon=lons, lat=lats, depth=np.repeat(5, n), time=np.repeat(start, n))
kernel = AdvectionRK4_3D
output_file = pset.ParticleFile(name=prefix + '_NEMO_3D.zarr', outputdt=timedelta(hours=1))
pset.execute(
    kernel, runtime=duration, dt=dt, output_file=output_file,
    recovery={ErrorCode.ErrorOutOfBounds: DeleteParticle},
)


# Visualize results

# In[7]:


get_ipython().run_cell_magic('capture', '', '\n# Make initial figure\nfig, ax = plt.subplots(figsize=(8, 8))\ncax = fig.add_axes([0.92, 0.15, 0.02, 0.7])\nl = ax.scatter([], [], s=50, c=[], vmin=0, vmax=10, edgecolor=\'k\') # "s" controls the size of the points\nt = ax.text(0.02, 0.02, \'\', transform=ax.transAxes)\ndata = xr.open_dataset(prefix + \'_NEMO_3D.zarr\') # This line opens the dataset of the particles, so you don\'t have to run the \n# particle simulation every time you come back to this notebook to play around with it, as long as you want to keep the particles the same\nax.contourf(gridlon, gridlat, tmask, levels=[-0.01, 0.01], colors=\'gray\')\nax.contour(gridlon, gridlat, tmask, levels=[-0.01, 0.01], colors=\'k\')\nax.set_xlim([-124.3, -123])\nax.set_ylim([48.7, 49.6])\nax.set_title(\'NEMO_3D\')\nax.set_aspect(1/np.sin(np.deg2rad(49)))\nfig.colorbar(l, cax=cax, label=\'Depth [m]\')\n\n# Init function\ndef init():\n    t.set_text(\'\')\n    l.set_offsets(np.empty((0, 2)))\n    l.set_array(np.empty(0))\n    return l, t,\n\n# Animate function\ndef animate(hour):\n    tstamp = data.time[0, hour].values.astype(\'datetime64[s]\').astype(datetime)\n    t.set_text(tstamp.strftime(\'%Y-%b-%d %H:%M UTC\'))\n    l.set_offsets(np.vstack([data.lon[:, hour], data.lat[:, hour]]).T)\n    l.set_array(data.z[:, hour])\n    return l, t,\n\n# Build animation\nanim = animation.FuncAnimation(fig, animate, init_func=init, frames=73, interval=100, blit=True)\n')


# In[8]:


# Render animation
HTML(anim.to_html5_video())


# This is the end of what I did using Ben's notebook example
# 
# From here on is what I have developed

# In[55]:


get_ipython().run_cell_magic('capture', '', "\n# Make initial vertical section profile figure\nfig, ax = plt.subplots(figsize=(8, 8)) # names the figure fig and sets the figure size\ncax = fig.add_axes([0.92, 0.15, 0.02, 0.7]) # adds colour bar and sets its position\nl = ax.scatter([data.lon], [data.z], s=50, c=[data.z], vmin=0, vmax=10, edgecolor='k')\nt = ax.text(0.02, 0.02, '', transform=ax.transAxes)\ndata = xr.open_dataset(prefix + '_NEMO_3D.zarr')\n#ax.contourf(gridlon, gridlat, tmask, levels=[-0.01, 0.01], colors='gray') # don't need these because it's a vertical plot\n#ax.contour(gridlon, gridlat, tmask, levels=[-0.01, 0.01], colors='k')\nax.set_xlim([-124.3, -123])\nax.set_ylim([100, 0])\nax.set_title('NEMO_3D')\nax.set_ylabel('Depth [m]')\n#ax.set_aspect(1/np.sin(np.deg2rad(49)))\nfig.colorbar(l, cax=cax, label='Depth [m]')\n\n# Init function\ndef init():\n    t.set_text('')\n    l.set_offsets(np.empty((0, 2)))\n    l.set_array(np.empty(0))\n    return l, t,\n\n# Animate function\ndef animate(hour):\n    tstamp = data.time[0, hour].values.astype('datetime64[s]').astype(datetime)\n    t.set_text(tstamp.strftime('%Y-%b-%d %H:%M UTC'))\n    l.set_offsets(np.vstack([data.lon[:, hour], data.z[:, hour]]).T)\n    l.set_array(data.z[:, hour])\n    return l, t,\n\n# Build animation\nanim = animation.FuncAnimation(fig, animate, init_func=init, frames=73, interval=100, blit=True)\n")


# In[56]:


# Render animation
HTML(anim.to_html5_video())


# In[27]:


data = xr.open_dataset(prefix + '_NEMO_3D.zarr')
print(data)


# In[33]:


print(data.lon[0, 0])


# Make side by side map and depth profile

# In[127]:


get_ipython().run_cell_magic('capture', '', "\n# Make initial figure\nfig, axs = plt.subplots(1, 2, figsize=(17, 8), gridspec_kw={'wspace': 0.1})\ncax = fig.add_axes([0.92, 0.15, 0.02, 0.7]) # do I need this in this side by side?\nl0 = axs[0].scatter([data.lon], [data.lat], s=50, c=[data.z], vmin=0, vmax=10, edgecolor='k')\nl1 = axs[1].scatter([data.lon], [data.z], s=50, c=[data.z], vmin=0, vmax=10, edgecolor='k')\nt0 = axs[0].text(0.02, 0.02, '', transform=axs[0].transAxes)\nt1 = axs[1].text(0.02, 0.02, '', transform=axs[1].transAxes)\ndata = xr.open_dataset(prefix + '_NEMO_3D.zarr') # This line opens the dataset of the particles, so you don't have to run the \n# particle simulation every time you come back to this notebook to play around with it, as long as you want to keep the particles the same\naxs[0].contourf(gridlon, gridlat, tmask, levels=[-0.01, 0.01], colors='gray')\naxs[0].contour(gridlon, gridlat, tmask, levels=[-0.01, 0.01], colors='k')\naxs[0].set_xlim([-124.4, -123])\naxs[1].set_xlim([-124.4, -123])\naxs[0].set_ylim([48.7, 49.8])\naxs[1].set_ylim([100, 0])\naxs[1].set_ylabel('Depth [m]')\naxs[0].set_aspect(1/np.sin(np.deg2rad(49)))\nfig.colorbar(l1, cax=cax, label='Depth [m]')\n\n# Init function\ndef init():\n    t0.set_text('')\n    t1.set_text('')\n    l0.set_offsets(np.empty((0, 2)))\n    l1.set_offsets(np.empty((0, 2)))\n    l0.set_array(np.empty(0))\n    l1.set_array(np.empty(0))\n    return l0, l1, t0, t1,\n\n# Animate function\ndef animate(hour):\n    tstamp = data.time[0, hour].values.astype('datetime64[s]').astype(datetime)\n    t0.set_text(tstamp.strftime('%Y-%b-%d %H:%M UTC'))\n    t1.set_text(tstamp.strftime('%Y-%b-%d %H:%M UTC'))\n    l0.set_offsets(np.vstack([data.lon[:, hour], data.lat[:, hour]]).T)\n    l1.set_offsets(np.vstack([data.lon[:, hour], data.z[:, hour]]).T)\n    l0.set_array(data.z[:, hour])\n    l1.set_array(data.z[:, hour])\n    return l0, l1, t0, t1,\n\n# Build animation\nanim = animation.FuncAnimation(fig, animate, init_func=init, frames=73, interval=100, blit=True)\n")


# In[128]:


# Render animation
HTML(anim.to_html5_video())


# In[76]:


pset.show(field=fieldset.W)


# In[51]:


fieldset.U.show()


# In[65]:


plotTrajectoriesFile(prefix + '_NEMO_3D.zarr')


# In[70]:


plotTrajectoriesFile(prefix + '_NEMO_3D.zarr', mode='movie2d_notebook')


# In[ ]:


plotTrajectoriesFile(prefix + '_NEMO_3D.zarr', mode='movie2d_notebook', recordedvar=data.z)


# Tests with different particle initializations

# In[5]:


# Paths and filenames
paths = {
    'coords': '/data/SalishSeaCast/grid/coordinates_seagrid_SalishSea201702.nc',
    'mask': '/data/SalishSeaCast/grid/mesh_mask201702.nc',
    'results': '/ocean/mattmiller/MOAD/analysis-matt/results/parcels/test',
}

# Load coords and mask files and extract grid variables
coords, mask = [xr.open_dataset(paths[key], decode_times=False) for key in ('coords', 'mask')]
gridlon, gridlat = [coords[key][0, ...].values for key in ('glamt', 'gphit')]
tmask = mask.tmask[0, 0, ...].values

# Define release parameters
location = 'Haro Strait' # there are predefined location names built into SalishSeaTools, under salishsea_tools.places.PLACES which can be found easily on github
# How would I set a custom lat lon for start location?
n = 100   # number of particles
r = 5   # radius of particle cloud [m]

# Start time, duration and timestep
start = datetime(2019, 1, 1, 12, 30, 0) # year, month, day, hour, minute, second
duration = timedelta(days=3) # the total duration of the simulation
dt = timedelta(seconds=90) # the timestep - toggle between + and - to run forwards or reverse
# is this the timestep that parcels uses or of the fieldset from SalishSeaCast?

# Create Gaussian distribution around release point
mean, cov = [0, 0], [[r**2, 0], [0, r**2]]
x_offset, y_offset = np.random.multivariate_normal(mean, cov, n).T
lon, lat = (-123.25, 48.5) # enter custom release coordinates here (rename it above as "location")
lons = lon + x_offset / 111000 / np.cos(np.deg2rad(lat)) # this code might be used to convert from SalishSeaCast's grid to a lat lon grid, since it is an angled domain grid
lats = lat + y_offset / 111000

# Forcing daterange (I add 1-day buffers)
daterange = [start - timedelta(days=1), start + duration + timedelta(days=1)]

# Output file prefix
strings = [location] + [t.strftime('%Y%m%dT%H%M%S') for t in (start, start + duration)]
prefix = os.path.join(paths['results'], '_'.join(strings))


# In[6]:


# Load SalishSeaCast results into fieldset
fieldset = fieldset_from_nemo(daterange, paths['coords'])


# In[7]:


# Execute NEMO-only, 3D run, release randomly dispersed over 60 m depth - originally tried 100 m and this location is apparently shallower so a bunch of particles immediately were stuck
pset = ParticleSet.from_list(fieldset, JITParticle, lon=lons, lat=lats, depth=np.random.random_sample(n)*(60), time=np.repeat(start, n))
kernel = AdvectionRK4_3D
output_file = pset.ParticleFile(name=prefix + '_NEMO_3D.zarr', outputdt=timedelta(hours=1))
pset.execute(
    kernel, runtime=duration, dt=dt, output_file=output_file,
    recovery={ErrorCode.ErrorOutOfBounds: DeleteParticle},
)


# In[10]:


get_ipython().run_cell_magic('capture', '', "\n# Make initial figure\nfig, axs = plt.subplots(1, 2, figsize=(17, 8), gridspec_kw={'wspace': 0.1})\ncax = fig.add_axes([0.92, 0.15, 0.02, 0.7]) # do I need this in this side by side?\ndata = xr.open_dataset(prefix + '_NEMO_3D.zarr') # This line opens the dataset of the particles, so you don't have to run the \n# particle simulation every time you come back to this notebook to play around with it, as long as you want to keep the particles the same\nl0 = axs[0].scatter([data.lon], [data.lat], s=50, c=[data.z], vmin=0, vmax=100, edgecolor='k')\nl1 = axs[1].scatter([data.lon], [data.z], s=50, c=[data.z], vmin=0, vmax=100, edgecolor='k')\nt0 = axs[0].text(0.02, 0.02, '', transform=axs[0].transAxes)\nt1 = axs[1].text(0.02, 0.02, '', transform=axs[1].transAxes)\naxs[0].contourf(gridlon, gridlat, tmask, levels=[-0.01, 0.01], colors='gray')\naxs[0].contour(gridlon, gridlat, tmask, levels=[-0.01, 0.01], colors='k')\naxs[0].set_xlim([-125, -122])\naxs[1].set_xlim([-125, -122])\naxs[0].set_ylim([47.5, 49.5])\naxs[1].set_ylim([200, 0])\naxs[1].set_ylabel('Depth [m]')\naxs[0].set_aspect(1/np.sin(np.deg2rad(49)))\nfig.colorbar(l1, cax=cax, label='Depth [m]')\n\n# Init function\ndef init():\n    t0.set_text('')\n    t1.set_text('')\n    l0.set_offsets(np.empty((0, 2)))\n    l1.set_offsets(np.empty((0, 2)))\n    l0.set_array(np.empty(0))\n    l1.set_array(np.empty(0))\n    return l0, l1, t0, t1,\n\n# Animate function\ndef animate(hour):\n    tstamp = data.time[0, hour].values.astype('datetime64[s]').astype(datetime)\n    t0.set_text(tstamp.strftime('%Y-%b-%d %H:%M UTC'))\n    t1.set_text(tstamp.strftime('%Y-%b-%d %H:%M UTC'))\n    l0.set_offsets(np.vstack([data.lon[:, hour], data.lat[:, hour]]).T)\n    l1.set_offsets(np.vstack([data.lon[:, hour], data.z[:, hour]]).T)\n    l0.set_array(data.z[:, hour])\n    l1.set_array(data.z[:, hour])\n    return l0, l1, t0, t1,\n\n# Build animation\nanim = animation.FuncAnimation(fig, animate, init_func=init, frames=73, interval=100, blit=True)\n")


# In[11]:


# Render animation
HTML(anim.to_html5_video())


# When originally set to release over 100 m, some of the particles were "stuck" right from initialization, which tells me they were initialized greater than the bottom depth at that location. I re-ran using 60 m max release depth and it fixed it but still looks like one gets stuck right away, at ~50m? Weird because deeper particles don't get stuck

# How to index certain values from the dataset: call a variable from the dataset e.g. data.lon and use square brackets: e.g. data.lon[:,0] calls for all (:) particles but just their entry at time = 0. data.lon[:,-1] calls all particles (:) but their position at the very last time in the array. data.lon[0,:] calls just the first particle (remember python indexing starts at 0) and its position at all time steps

# In[16]:


# Make initial figure
fig, axs = plt.subplots(1, 2, figsize=(17, 8), gridspec_kw={'wspace': 0.1})
cax = fig.add_axes([0.92, 0.15, 0.02, 0.7]) # do I need this in this side by side?
l0 = axs[0].scatter([data.lon[:,-1]], [data.lat[:,-1]], s=50, c=[data.z[:,-1]], vmin=0, vmax=100, edgecolor='k')
l1 = axs[1].scatter([data.lon[:,-1]], [data.z[:,-1]], s=50, c=[data.z[:,-1]], vmin=0, vmax=100, edgecolor='k')
t0 = axs[0].text(0.02, 0.02, '', transform=axs[0].transAxes)
t1 = axs[1].text(0.02, 0.02, '', transform=axs[1].transAxes)
data = xr.open_dataset(prefix + '_NEMO_3D.zarr') # This line opens the dataset of the particles, so you don't have to run the 
# particle simulation every time you come back to this notebook to play around with it, as long as you want to keep the particles the same
axs[0].contourf(gridlon, gridlat, tmask, levels=[-0.01, 0.01], colors='gray')
axs[0].contour(gridlon, gridlat, tmask, levels=[-0.01, 0.01], colors='k')
axs[0].set_xlim([-125, -122])
axs[1].set_xlim([-125, -122])
axs[0].set_ylim([47.5, 49.5])
axs[1].set_ylim([200, 0])
axs[1].set_ylabel('Depth [m]')
axs[0].set_aspect(1/np.sin(np.deg2rad(49)))
fig.colorbar(l1, cax=cax, label='Depth [m]')