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

# 2D Parcels Test Model (depth and lon) - Matt

# # Load packages and functions

# In[ ]:


get_ipython().run_line_magic('matplotlib', 'inline')
from parcels import Field, FieldSet, ParticleSet, Variable, JITParticle, ErrorCode
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from matplotlib import rc, animation
import xarray as xr
from datetime import datetime
import os

from IPython.display import Image
rc('animation', html='html5')


# # Build parcels particle simulation

# ## Step 1: Define parcels FieldSet

# In[ ]:


# Define SalishSeaCast results filepath
def make_prefix(date, res='h'):
    """Construct path prefix for local SalishSeaCast results given date object and paths dict
    e.g., /results2/SalishSea/nowcast-green.201905/daymonthyear/SalishSea_1h_yyyymmdd_yyyymmdd
    """
    path = '/results2/SalishSea/nowcast-green.202111/'
    datestr = '_'.join(np.repeat(date.strftime('%Y%m%d'), 2))
    folder = date.strftime("%d%b%y").lower()
    prefix = os.path.join(path, f'{folder}/SalishSea_1{res}_{datestr}')
    
    return prefix


# In[ ]:


# Pick hour, day and location for simulation
hour = 0 # Choose start hour
path_NEMO = make_prefix(datetime(2018, 8, 1)) # Choose start year, month, day
lati = 49.224563 # Choose coords
loni = -123.605357
jjii = xr.open_dataset('/home/sallen/MEOPAR/grid/grid_from_lat_lon_mask999.nc') # This is a file that gives the closest SalishSeaCast grid point (i and j) for a given lat and lon 
j = [jjii.jj.sel(lats=lati, lons=loni, method='nearest').item()][0] # j is the index along the SalishSeaCast grid (think y axis of model domain, or similar to latitude), while i is the index across the grid
i = [jjii.ii.sel(lats=lati, lons=loni, method='nearest').item()][0]

# Specify some simulation parameters
T = 6 * 86400 #s (run time) (86400 s = 24 hours/1 day) - initially 5e3 * 100 (this is about 5.7 days)
dt = 5 #s (timestep)
N = 10e3 #number of particles
outputdt = 3600 #s (3600 = 1 hr) (how often do you want output?) initially 500
outputpath = '/ocean/mattmiller/MOAD/results/parcels/test/Outputmix_sink_and_swim_6_days.zarr' # this makes sure the results are saved in my /results folder located beside /analysis-matt


# In[ ]:


# Define domain, velocity fields and Kz 
dim = 100
dep = len(depth)
lon = np.linspace(0., 2e3, dim, dtype=np.float32)

# Define variables for parcels FieldSet
U = Field('U', np.zeros((dep, dim), dtype=np.float32), lon=lon, depth=depth) # in parcels, 'U' represents the zonal flow velocity (zonal = east-to-west/west-to-east)
V = Field('V', np.zeros((dep, dim), dtype=np.float32), lon=lon, depth=depth) # in parcels, 'V' represents the meridional flow velocity (meridional = north-to-south/south-to-north)
Kz_data = np.zeros((dep, dim), dtype=np.float32)
for i in range(dim):
    Kz_data[:,i]=Kz_col   
Kz = Field('Kz', Kz_data, grid=U.grid)

# Build parcels FieldSet
fieldset = FieldSet(U,V)
fieldset.add_field(Kz)


# ## Step 2: Define parcels ParticleSet

# additional Variables can be added to the particles (e.g. temperature, to keep track of the temperature that particles experience).

# In[ ]:


# Define a new particleclass with Variable 'age' with initial value 0.
AgeParticle = parcels.JITParticle.add_variable(parcels.Variable("age", initial=0))

pset = parcels.ParticleSet(
    fieldset=fieldset,  # the fields that the particleset uses
    pclass=AgeParticle,  # define the type of particle
    lon=29,  # release longitude
    lat=-33,  # release latitude
)


# ## Step 3: Define parcels particle kernels

# In[ ]:


# Create a custom kernel which displaces each particle southward
def NorthVel(particle, fieldset, time):
    if time > 10 * 86400 and time < 10.2 * 86400:
        vvel = -1e-4
        particle_dlat += vvel * particle.dt


# Create a custom kernel which keeps track of the particle age (minutes)
def Age(particle, fieldset, time):
    particle.age += particle.dt / 3600


# define all kernels to be executed on particles using an (ordered) list
kernels = [Age, NorthVel, parcels.AdvectionRK4]


# ## Step 4: Execute parcels simulation

# In[ ]:


output_file = pset.ParticleFile(
    name = "Outputmix_sink_and_swim.zarr",  # the name of the output file
    outputdt = 3600,  # the time period between consecutive output steps
    chunks = (1, 10),  # the chunking of the output file (number of particles, timesteps)
)

pset.execute(
    kernels,  # the list of kernels (which defines how particles move)
    runtime = 86400 * 24,  # the total length of the run in seconds
    dt = 300,  # the timestep of the kernel in seconds
    output_file = output_file,
)


# ## Load results from previous parcels simulations

# In[ ]:


dat = xr.load_dataset('/ocean/mattmiller/MOAD/results/parcels/test/Outputmix_sink_and_swim.zarr')


# In[ ]:


# View data
print(dat)