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

# Plot Antarctic Tidal Currents
# =============================
# 
# Demonstrates plotting hourly tidal currents around Antarctica
# 
# OTIS format tidal solutions provided by Ohio State University and ESR  
# - http://volkov.oce.orst.edu/tides/region.html  
# - https://www.esr.org/research/polar-tide-models/list-of-polar-tide-models/
# - ftp://ftp.esr.org/pub/datasets/tmd/  
# 
# Finite Element Solution (FES) provided by AVISO  
# - https://www.aviso.altimetry.fr/en/data/products/auxiliary-products/global-tide-fes.html
# 
# #### Python Dependencies
#  - [numpy: Scientific Computing Tools For Python](https://www.numpy.org)  
#  - [scipy: Scientific Tools for Python](https://www.scipy.org/)  
#  - [pyproj: Python interface to PROJ library](https://pypi.org/project/pyproj/)  
#  - [netCDF4: Python interface to the netCDF C library](https://unidata.github.io/netcdf4-python/)  
#  - [matplotlib: Python 2D plotting library](http://matplotlib.org/)  
#  - [cartopy: Python package designed for geospatial data processing](https://scitools.org.uk/cartopy/docs/latest/)  
# 
# #### Program Dependencies
# 
# - `calc_astrol_longitudes.py`: computes the basic astronomical mean longitudes  
# - `convert_ll_xy.py`: convert lat/lon points to and from projected coordinates  
# - `load_constituent.py`: loads parameters for a given tidal constituent  
# - `load_nodal_corrections.py`: load the nodal corrections for tidal constituents  
# - `io.model.py`: retrieves tide model parameters for named tide models
# - `io.OTIS.py`: extract tidal harmonic constants from OTIS tide models  
# - `io.ATLAS.py`: extract tidal harmonic constants from ATLAS netcdf models  
# - `io.FES.py`: extract tidal harmonic constants from FES tide models  
# - `predict.py`: predict tidal values using harmonic constants  
# - `time.py`: utilities for calculating time operations
# - `utilities.py`: download and management utilities for files
# 
# This notebook uses Jupyter widgets to set parameters for calculating the tidal maps.  

# #### Load modules

# In[ ]:


import os
import pyproj
import datetime
import numpy as np
import matplotlib
matplotlib.rcParams['axes.linewidth'] = 2.0
matplotlib.rcParams["animation.html"] = "jshtml"
import matplotlib.pyplot as plt
import matplotlib.animation as animation
import cartopy.crs as ccrs
import ipywidgets as widgets
from IPython.display import HTML

# import tide programs
import pyTMD.io
import pyTMD.time
import pyTMD.predict
import pyTMD.tools
import pyTMD.utilities

# autoreload
get_ipython().run_line_magic('load_ext', 'autoreload')
get_ipython().run_line_magic('autoreload', '2')


# #### Set parameters  for program
# 
# - Model directory  
# - Tide model  
# - Date to run  

# In[ ]:


# available model list
model_list = sorted(pyTMD.io.model.global_current() + pyTMD.io.model.antarctic_current())
# display widgets for setting directory and model
TMDwidgets = pyTMD.tools.widgets()
TMDwidgets.model.options = model_list
TMDwidgets.model.value = 'CATS2008'
widgets.VBox([
    TMDwidgets.directory,
    TMDwidgets.model,
    TMDwidgets.atlas,
    TMDwidgets.compress,
    TMDwidgets.datepick
])


# #### Setup tide model parameters

# In[ ]:


# get model parameters
model = pyTMD.io.model(TMDwidgets.directory.value,
    format=TMDwidgets.atlas.value,
    compressed=TMDwidgets.compress.value
   ).current(TMDwidgets.model.value)


# #### Setup coordinates for calculating tidal currents

# In[ ]:


# create an image around Antarctica
xlimits = [-560.*5e3,560.*5e3]
ylimits = [-560.*5e3,560.*5e3]
spacing = [20e3,-20e3]
# x and y coordinates
x = np.arange(xlimits[0],xlimits[1]+spacing[0],spacing[0])
y = np.arange(ylimits[1],ylimits[0]+spacing[1],spacing[1])
xgrid,ygrid = np.meshgrid(x,y)
# x and y dimensions
nx = int((xlimits[1]-xlimits[0])/spacing[0])+1
ny = int((ylimits[0]-ylimits[1])/spacing[1])+1
# convert image coordinates from polar stereographic to latitude/longitude
crs1 = pyproj.CRS.from_epsg(3031)
crs2 = pyproj.CRS.from_epsg(4326)
transformer = pyproj.Transformer.from_crs(crs1, crs2, always_xy=True)
lon,lat = transformer.transform(xgrid.flatten(), ygrid.flatten())


# #### Calculate tide map

# In[ ]:


# convert from calendar date to days relative to Jan 1, 1992 (48622 MJD)
YMD = TMDwidgets.datepick.value
tide_time = pyTMD.time.convert_calendar_dates(YMD.year, YMD.month,
    YMD.day, hour=np.arange(24))
# delta time (TT - UT1) file
delta_file = pyTMD.utilities.get_data_path(['data','merged_deltat.data'])

# save tide currents
tide = {}
# iterate over u and v currents
for TYPE in model.type:
    # read tidal constants and interpolate to grid points
    if model.format in ('OTIS','ATLAS','ESR'):
        amp,ph,D,c = pyTMD.io.OTIS.extract_constants(lon, lat, model.grid_file,
            model.model_file['u'], model.projection, type=TYPE,
            method='spline', grid=model.format)
        DELTAT = np.zeros_like(tide_time)
    elif (model.format == 'netcdf'):
        amp,ph,D,c = pyTMD.io.ATLAS.extract_constants(lon, lat, model.grid_file,
            model.model_file[TYPE], type=TYPE, method='spline',
            scale=model.scale, compressed=model.compressed)
        DELTAT = np.zeros_like(tide_time)
    elif (model.format == 'GOT'):
        amp,ph,c = pyTMD.io.GOT.extract_constants(lon, lat, model.model_file[TYPE],
            method='spline', scale=model.scale,
            compressed=model.compressed)
        # interpolate delta times from calendar dates to tide time
        DELTAT = pyTMD.time.interpolate_delta_time(delta_file, tide_time)
    elif (model.format == 'FES'):
        amp,ph = pyTMD.io.FES.extract_constants(lon, lat, model.model_file[TYPE],
            type=TYPE, version=model.version, method='spline',
            scale=model.scale, compressed=model.compressed)
        c = model.constituents
        # interpolate delta times from calendar dates to tide time
        DELTAT = pyTMD.time.interpolate_delta_time(delta_file, tide_time)
        
    # calculate complex phase in radians for Euler's
    cph = -1j*ph*np.pi/180.0
    # calculate constituent oscillation
    hc = amp*np.exp(cph)

    # allocate for tide current map calculated every hour
    tide[TYPE] = np.ma.zeros((ny,nx,24))
    for hour in range(24):
        # predict tidal elevations at time and infer minor corrections
        TIDE = pyTMD.predict.map(tide_time[hour], hc, c, deltat=DELTAT[hour],
            corrections=model.format)
        MINOR = pyTMD.predict.infer_minor(tide_time[hour], hc, c,
            deltat=DELTAT[hour], corrections=model.format)
        # add major and minor components and reform grid
        tide[TYPE][:,:,hour] = np.reshape((TIDE+MINOR),(ny,nx))


# #### Create animation of hourly tidal oscillation

# In[ ]:


# output Antarctic Tidal Current Animation
projection = ccrs.Stereographic(central_longitude=0.0,
    central_latitude=-90.0,true_scale_latitude=-71.0)
# figure axis and image objects
ax1,im = ({},{})
fig, (ax1['u'],ax1['v']) = plt.subplots(num=1, ncols=2,
    figsize=(11.5,7), subplot_kw=dict(projection=projection))
vmin = np.min([tide['u'].min(),tide['v'].min()])
vmax = np.max([tide['u'].max(),tide['v'].max()])
extent = (xlimits[0],xlimits[1],ylimits[0],ylimits[1])
for TYPE,ax in ax1.items():
    # plot tidal currents
    im[TYPE] = ax.imshow(np.zeros((ny,nx)),
        interpolation='nearest', vmin=vmin, vmax=vmax,
        transform=projection, extent=extent, origin='upper',
        animated=True)
    # add high resolution cartopy coastlines
    ax.coastlines('10m')
    # set x and y limits
    ax.set_xlim(xlimits)
    ax.set_ylim(ylimits)
    # stronger linewidth on frame
    ax.spines['geo'].set_linewidth(2.0)
    ax.spines['geo'].set_capstyle('projecting')

# Add colorbar with a colorbar axis
# Add an axes at position rect [left, bottom, width, height]
cbar_ax = fig.add_axes([0.085, 0.075, 0.83, 0.035])
# extend = add extension triangles to upper and lower bounds
# options: neither, both, min, max
cbar = fig.colorbar(im['u'], cax=cbar_ax, extend='both',
    extendfrac=0.0375, drawedges=False, orientation='horizontal')
# rasterized colorbar to remove lines
cbar.solids.set_rasterized(True)
# Add label to the colorbar
cbar.ax.set_xlabel(f'{model.name} Tidal Velocity', fontsize=13)
cbar.ax.set_ylabel('cm/s', fontsize=13, rotation=0)
cbar.ax.yaxis.set_label_coords(1.04, 0.15)
# ticks lines all the way across
cbar.ax.tick_params(which='both', width=1, length=18,
    labelsize=13, direction='in')

# add title (date and time)
ttl = fig.suptitle(None, fontsize=13)
# adjust subplot within figure
fig.subplots_adjust(left=0.02,right=0.98,bottom=0.1,top=0.98,wspace=0.04)
           
# animate each map
def animate_maps(hour):
    # set map data iterating over u and v currents
    for TYPE in model.type:
        im[TYPE].set_data(tide[TYPE][:,:,hour])
    # set title
    args = (YMD.year,YMD.month,YMD.day,hour)
    ttl.set_text('{0:4d}-{1:02d}-{2:02d}T{3:02d}:00:00'.format(*args))

# set animation
anim = animation.FuncAnimation(fig, animate_maps, frames=24)
get_ipython().run_line_magic('matplotlib', 'inline')
HTML(anim.to_jshtml())


#