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

# # Frist real run POP Data

# In[1]:


from parcels import FieldSet, ParticleSet, Variable, JITParticle, AdvectionRK4_3D, plotTrajectoriesFile,ErrorCode
import numpy as np
import math
import datetime
from datetime import date
from datetime import timedelta as delta
#from datetime import datetime
from operator import attrgetter
import netCDF4 as nc
import matplotlib.pyplot as plt
import cartopy.crs as ccrs
from glob import glob
from parcels import version
import netCDF4
import numpy.matlib
import time
import os
print('Parcels version: ', version)


# ### Prepare loading data

# In[2]:


###### NEW LOADING CELL ######
################################################################################################################################
## loading data - creating field set
# load all data
#data_path = '/home/cakatsman/NoraFried/IrmSea_data_POP/'
data_path = '/data/nora/POP/'
files_all = sorted(glob(data_path+'t.t0.*.nc'))

# select which data to use e.g. only use summer to summer
start_day = 182
files_wanted = files_all[start_day:start_day+365]

# put all files together
#files = np.append(files_wanted,files_wanted,files_wanted,files_wanted,files_wanted)

files = files_wanted
years2run = 1
for x in range(years2run-1):
    files = np.append(files, files_wanted)
    
print('==== start July year 1 until June year 2 ====')


# In[3]:


# additional definitions for loading files
filenames = { 'U': {'lon':files[0], 'lat' : files[0], 'depth' : files[0], 'data' : files},
              'V': {'lon':files[0], 'lat' : files[0], 'depth' : files[0], 'data' : files},
              'W': {'lon':files[0], 'lat' : files[0], 'depth' : files[0], 'data' : files}}

variables = {'U': 'UVEL', 
             'V': 'VVEL',
             'W': 'WVEL'}

dimensions =   {'U':{'lon': 'U_LON_2D', 'lat': 'U_LAT_2D','depth' : 'w_dep'},
                'V':{'lon': 'U_LON_2D', 'lat': 'U_LAT_2D','depth' : 'w_dep'},
                'W':{'lon': 'U_LON_2D', 'lat': 'U_LAT_2D','depth' : 'w_dep'}}  


# ### Create time stamps

# In[4]:


## TIMESTAMPS 
start_dt = datetime.datetime.now()

timestamps_prep = []

for i in range(1,len(files)+1):
 timestamps_prep.append((start_dt+datetime.timedelta(days=i)).strftime('%Y-%m-%d'))

#print(timestamps_prep)

timestamps = np.expand_dims(np.array([np.datetime64('%s' %m) for m in timestamps_prep]), axis=1)

#print(np.shape(timestamps))

time_cycle = 365*5
print("total run time:", time_cycle,"days = ",time_cycle/365,"years")


# ### Create fieldset

# In[5]:


#fieldset = FieldSet.from_pop(filenames, variables, dimensions, timestamps=timestamps, depth_units = 'cm')
fieldset = FieldSet.from_pop(filenames, variables, dimensions, timestamps=timestamps, time_periodic=datetime.timedelta(days=time_cycle),field_chunksize=False)


# ### 2. Particle Set

# In[6]:


## LOAD OSNAP coordinates
import scipy.io as sio
os.getcwd()

OSNAPlatlon = sio.loadmat('/data/nora/OSNAPlatlon.mat')

sorted(OSNAPlatlon.keys())

np.shape(OSNAPlatlon['osnap'])

#print(OSNAPlatlon['osnap'])
osnap_latlon = OSNAPlatlon['osnap']
osnap_lat = osnap_latlon[:,0] 
osnap_lon = osnap_latlon[:,1]


#print(osnap_lat[23],osnap_lon[23]) # Western most OSNAP East mooring
#print(osnap_lat[46],osnap_lon[46]) # M1

#indices_osnap_east = [24:33 41:55]

west_coor = [osnap_lat[23],osnap_lon[23]]
east_coor = [osnap_lat[46],osnap_lon[46]]


# In[7]:


# start particles along a line at one latitude
npart = 120                                 # number of particles
#latp = 59 * np.ones(npart) 
#lonpp = [i for i in np.linspace(-36,-29,npart)] # longitude for particles

# deploy along one latitude
latp = [i for i in np.linspace(west_coor[0],east_coor[0],npart)]    # deploy along one latitude
lonp = [i for i in np.linspace(west_coor[1],east_coor[1],npart)] # longitude for particles
#print(lonp,latp)
repeatp= delta(days=365)                        # Optional interval (in seconds) on which to 
                                               # repeat the release of the ParticleSet
min_depth = 1050
max_depth = 1500 #2000
depth_levels = 10 #101
#depthp = [i for i in np.linspace(0,max_depth,npart)] 
depthpt = np.transpose(np.tile(np.linspace(min_depth,max_depth,depth_levels),(npart,1)))
latt = np.tile(latp,(depth_levels,1))                                       # make matrix for depth
lont = np.tile(lonp,(depth_levels,1))

#print(np.shape(depthpt))

#print(np.shape(latt))

print("number of particles:", npart)
print("maximum depth:",max_depth,"meters")
print("number of layers:",depth_levels," layerthickness:",(max_depth-min_depth)/(depth_levels-1),"meters")


# In[8]:


# extend the array for time dimension
time_len= 365
lat_parcels = np.tile(latt,(time_len,1,1)) 
lon_parcels = np.tile(lont,(time_len,1,1)) 
depth_parcels = np.tile(depthpt,(time_len,1,1)) 
time_parcels = np.transpose(np.tile(timestamps[0:365],(depth_levels,1,npart)),(1,0,2))


#print(np.shape(timestamps))
print("size lat:",np.shape(lat_parcels))
print("size depth:", np.shape(depth_parcels))
print("size time:", np.shape(time_parcels))


# In[9]:


print('particle class: JITParticle')
print("longitude west - east:",lon_parcels[1,1,0:2])
print("latitude west - east:",lat_parcels[1,1,0:2])


# In[10]:


pset = ParticleSet.from_list(fieldset=fieldset, 
                             pclass=JITParticle,    # particle type
                             lon = lon_parcels,     #lont,            
                             lat = lat_parcels,     #latt,            
                             depth = depth_parcels, #depthpt,  
                             time = time_parcels)#,
                             #repeatdt = repeatp)
#print('NO REPEAT ANYMORE')                            


# ### Kernels

# In[11]:


# decide which kernels to use
# --> AdvectionRK4_3D
    
# define all kernels to be executed on particles
kernels = pset.Kernel(AdvectionRK4_3D) #+ pset.Kernel(periodicBC)

def SubmergeParticle(particle, fieldset, time):
    particle.depth = 0.01  # Minimum depth is 0.5m
    particle.time = time + particle.dt  # to not trigger kernels again, otherwise infinite loop

def DeleteParticle(particle, fieldset, time):
    particle.delete()


# In[12]:


#pset.show()


# ### 4. Execution and output

# ### 4.2 backward

# In[15]:


output_file = pset.ParticleFile(name="parcels_backward_120part_5years_1050m_1500m.nc", outputdt=delta(days=2)) # file name and time step of output


start_time = time.time()

pset.execute(kernels,                   # the kernel which define how particles move
             runtime = delta(days=365*5),  # the total length of the run 
             dt = delta(hours=-1),        # BACKWARD
             recovery={ErrorCode.ErrorOutOfBounds: DeleteParticle,
                       ErrorCode.ErrorThroughSurface: SubmergeParticle},
             output_file=output_file)

output_file.export()
output_file.close()

print("--- %s minutes ---" % np.round(((time.time() - start_time)/60)))


# ### 5. Plotting

# In[ ]: