%matplotlib inline
from datetime import timedelta, datetime
import xarray as xr
import numpy as np
import matplotlib.pyplot as plt
import matplotlib

from matplotlib.animation import FuncAnimation
import matplotlib.animation as animation

from IPython.display import HTML

import cartopy
import cartopy.crs as ccrs
from cartopy.mpl.gridliner import LONGITUDE_FORMATTER, LATITUDE_FORMATTER
import matplotlib.ticker as mticker
import cartopy.feature as cfeature

rootdir = "/Users/Gomez023/src/git/"

filedir = "parcels_gallery/data/"

ds_subregion = xr.open_dataset(rootdir + filedir + 'Particle_AZO_grid100000p_wtides_0601_hourly_MONTH_subregion.nc')

outputdt = timedelta(hours=1)
timerange = np.arange(np.nanmin(ds_subregion['time'].values),
                      (np.datetime64('2010-06-29T00:30:00.000000000'))+np.timedelta64(outputdt), 
                      outputdt)  # timerange in nanoseconds

%%capture

fig = plt.figure(figsize=(12, 8))#, constrained_layout=True)

ax1 = plt.subplot(111, projection=ccrs.PlateCarree())

ax1.set_xlim([-32, -18])
ax1.set_ylim([30, 40])

resol = '10m'  # use data at this scale
land = cartopy.feature.NaturalEarthFeature('physical', 'land', \
    scale=resol, edgecolor='k', facecolor=cfeature.COLORS['land'])
ocean = cartopy.feature.NaturalEarthFeature('physical', 'ocean', \
    scale=resol, edgecolor='none', facecolor=cfeature.COLORS['water'])

ax1.add_feature(land, facecolor='burlywood', zorder=50) #beige
ax1.add_feature(ocean, linewidth=0.2 )

gl = ax1.gridlines(crs=ccrs.PlateCarree(), draw_labels=True,
                  linewidth=2, color='gray', alpha=0.5, linestyle='-')
gl.xlabels_top = False
gl.ylabels_right = False
gl.xlines = False
gl.ylines = False
gl.xlocator = mticker.FixedLocator([-35, -30., -25., -20.])
gl.ylocator = mticker.FixedLocator([30., 33., 36., 39.]) #np.arange(30., 42, 3)
gl.xformatter = LONGITUDE_FORMATTER
gl.yformatter = LATITUDE_FORMATTER
gl.ylabel_style = {'size': 12}#, 'color': 'gray'}
gl.xlabel_style = {'size': 12}


time_id = np.where(ds_subregion['time'] == timerange[0]) # Indices of the data where time = 0

scatter_subregion = ax1.scatter(ds_subregion['lon'].values[time_id], ds_subregion['lat'].values[time_id], s=1, c='r', transform=ccrs.PlateCarree(), zorder=20)

t = str(timerange[0]) 
title = plt.suptitle('Particles at t = ' + str(t[0:10]) + ' ' + str(t[11:16]), size=24)

def animate(i):
    t = str(timerange[i])
    title.set_text('Particles at t = ' + str(t[0:10]) + ' ' + str(t[11:16]))
    
    time_id = np.where(ds_subregion['time'] == timerange[i])
    scatter_subregion.set_offsets(np.c_[ds_subregion['lon'].values[time_id], ds_subregion['lat'].values[time_id]])

ax1.set_title('Tidal', size=20)

anim = FuncAnimation(fig, animate, frames = len(timerange), interval=500)

# Set up formatting for the movie files
Writer = animation.writers['ffmpeg']
writer = Writer(fps=15, metadata=dict(artist='Me'), bitrate=1800)

anim.save(rootdir + 'parcels_gallery/images/anim_June_s1_TIDAL_all_subregion.mp4', writer=writer, dpi=500)