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

# Plot Ross Ice Shelf Map
# =======================
# 
# Demonstrates plotting hourly tidal displacements for the Ross Ice Shelf
# 
# 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/  
# 
# Global Tide Model (GOT) solutions provided by Richard Ray at GSFC  
# 
# Finite Element Solution (FES) provided by AVISO  
# - https://www.aviso.altimetry.fr/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  
# - `calc_delta_time.py`: calculates difference between universal and dynamic time  
# - `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  
# - `infer_minor_corrections.py`: return corrections for minor constituents  
# - `read_tide_model.py`: extract tidal harmonic constants from OTIS tide models  
# - `read_netcdf_model.py`: extract tidal harmonic constants from netcdf models  
# - `read_GOT_model.py`: extract tidal harmonic constants from GSFC GOT models  
# - `read_FES_model.py`: extract tidal harmonic constants from FES tide models  
# - `predict_tide.py`: predict tidal elevation at a single time using harmonic constants  
# 
# This notebook uses Jupyter widgets to set parameters for calculating the tidal maps.  
# The widgets can be installed as described below.  
# ```
# pip3 install --user ipywidgets
# jupyter nbextension install --user --py widgetsnbextension
# jupyter nbextension enable --user --py widgetsnbextension
# jupyter-notebook
# ```

# #### 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 pyTMD.time
from pyTMD.calc_delta_time import calc_delta_time
from pyTMD.read_tide_model import extract_tidal_constants
from pyTMD.read_netcdf_model import extract_netcdf_constants
from pyTMD.read_GOT_model import extract_GOT_constants
from pyTMD.read_FES_model import extract_FES_constants
from pyTMD.infer_minor_corrections import infer_minor_corrections
from pyTMD.predict_tide import predict_tide
#-- 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[ ]:


#-- set the directory with tide models
dirText = widgets.Text(
    value=os.getcwd(),
    description='Directory:',
    disabled=False
)

#-- dropdown menu for setting tide model
model_list = ['CATS0201','CATS2008','TPXO9-atlas','TPXO9-atlas-v2',
    'TPXO9-atlas-v3','TPXO9.1','TPXO8-atlas','TPXO7.2',
    'GOT4.7','GOT4.8','GOT4.10','FES2014']
modelDropdown = widgets.Dropdown(
    options=model_list,
    value='CATS2008',
    description='Model:',
    disabled=False,
)

#-- date picker widget for setting time
datepick = widgets.DatePicker(
    description='Date:',
    value = datetime.date.today(),
    disabled=False
)

#-- display widgets for setting directory, model and date
widgets.VBox([dirText,modelDropdown,datepick])


# #### Setup tide model parameters

# In[ ]:


#-- directory with tide models
tide_dir = os.path.expanduser(dirText.value)
MODEL = modelDropdown.value
if (MODEL == 'CATS0201'):
    grid_file = os.path.join(tide_dir,'cats0201_tmd','grid_CATS')
    model_file = os.path.join(tide_dir,'cats0201_tmd','h0_CATS02_01')
    model_format = 'OTIS'
    EPSG = '4326'
    TYPE = 'z'
elif (MODEL == 'CATS2008'):
    grid_file = os.path.join(tide_dir,'CATS2008','grid_CATS2008')
    model_file = os.path.join(tide_dir,'CATS2008','hf.CATS2008.out')
    model_format = 'OTIS'
    EPSG = 'CATS2008'
    TYPE = 'z'
elif (MODEL == 'TPXO9-atlas'):
    model_directory = os.path.join(tide_dir,'TPXO9_atlas')
    grid_file = 'grid_tpxo9_atlas.nc.gz'
    model_files = ['h_q1_tpxo9_atlas_30.nc.gz','h_o1_tpxo9_atlas_30.nc.gz',
        'h_p1_tpxo9_atlas_30.nc.gz','h_k1_tpxo9_atlas_30.nc.gz',
        'h_n2_tpxo9_atlas_30.nc.gz','h_m2_tpxo9_atlas_30.nc.gz',
        'h_s2_tpxo9_atlas_30.nc.gz','h_k2_tpxo9_atlas_30.nc.gz',
        'h_m4_tpxo9_atlas_30.nc.gz','h_ms4_tpxo9_atlas_30.nc.gz',
        'h_mn4_tpxo9_atlas_30.nc.gz','h_2n2_tpxo9_atlas_30.nc.gz']
    model_format = 'netcdf'
    TYPE = 'z'
    SCALE = 1.0/1000.0
elif (MODEL == 'TPXO9-atlas-v2'):
    model_directory = os.path.join(tide_dir,'TPXO9_atlas_v2')
    grid_file = 'grid_tpxo9_atlas_30_v2.nc.gz'
    model_files = ['h_q1_tpxo9_atlas_30_v2.nc.gz','h_o1_tpxo9_atlas_30_v2.nc.gz',
        'h_p1_tpxo9_atlas_30_v2.nc.gz','h_k1_tpxo9_atlas_30_v2.nc.gz',
        'h_n2_tpxo9_atlas_30_v2.nc.gz','h_m2_tpxo9_atlas_30_v2.nc.gz',
        'h_s2_tpxo9_atlas_30_v2.nc.gz','h_k2_tpxo9_atlas_30_v2.nc.gz',
        'h_m4_tpxo9_atlas_30_v2.nc.gz','h_ms4_tpxo9_atlas_30_v2.nc.gz',
        'h_mn4_tpxo9_atlas_30_v2.nc.gz','h_2n2_tpxo9_atlas_30_v2.nc.gz']
    model_format = 'netcdf'
    TYPE = 'z'
    SCALE = 1.0/1000.0
elif (MODEL == 'TPXO9.1'):
    grid_file = os.path.join(tide_dir,'TPXO9.1','DATA','grid_tpxo9')
    model_file = os.path.join(tide_dir,'TPXO9.1','DATA','h_tpxo9.v1')
    model_format = 'OTIS'
    EPSG = '4326'
    TYPE = 'z'
elif (MODEL == 'TPXO8-atlas'):
    grid_file = os.path.join(tide_dir,'tpxo8_atlas','grid_tpxo8atlas_30_v1')
    model_file = os.path.join(tide_dir,'tpxo8_atlas','hf.tpxo8_atlas_30_v1')
    model_format = 'ATLAS'
    EPSG = '4326'
    TYPE = 'z'
elif (MODEL == 'TPXO7.2'):
    grid_file = os.path.join(tide_dir,'TPXO7.2_tmd','grid_tpxo7.2')
    model_file = os.path.join(tide_dir,'TPXO7.2_tmd','h_tpxo7.2')
    model_format = 'OTIS'
    EPSG = '4326'
    TYPE = 'z'
elif (MODEL == 'GOT4.7'):
    model_directory = os.path.join(tide_dir,'GOT4.7','grids_oceantide')
    model_files = ['q1.d.gz','o1.d.gz','p1.d.gz','k1.d.gz','n2.d.gz',
        'm2.d.gz','s2.d.gz','k2.d.gz','s1.d.gz','m4.d.gz']
    c = ['q1','o1','p1','k1','n2','m2','s2','k2','s1','m4']
    model_format = 'GOT'
    SCALE = 1.0/100.0
elif (MODEL == 'GOT4.8'):
    model_directory = os.path.join(tide_dir,'got4.8','grids_oceantide')
    model_files = ['q1.d.gz','o1.d.gz','p1.d.gz','k1.d.gz','n2.d.gz',
        'm2.d.gz','s2.d.gz','k2.d.gz','s1.d.gz','m4.d.gz']
    c = ['q1','o1','p1','k1','n2','m2','s2','k2','s1','m4']
    model_format = 'GOT'
    SCALE = 1.0/100.0
elif (MODEL == 'GOT4.10'):
    model_directory = os.path.join(tide_dir,'GOT4.10c','grids_oceantide')
    model_files = ['q1.d.gz','o1.d.gz','p1.d.gz','k1.d.gz','n2.d.gz',
        'm2.d.gz','s2.d.gz','k2.d.gz','s1.d.gz','m4.d.gz']
    c = ['q1','o1','p1','k1','n2','m2','s2','k2','s1','m4']
    model_format = 'GOT'
    SCALE = 1.0/100.0
elif (MODEL == 'FES2014'):
    model_directory = os.path.join(tide_dir,'fes2014','ocean_tide')
    model_files = ['2n2.nc.gz','eps2.nc.gz','j1.nc.gz','k1.nc.gz',
        'k2.nc.gz','l2.nc.gz','la2.nc.gz','m2.nc.gz','m3.nc.gz','m4.nc.gz',
        'm6.nc.gz','m8.nc.gz','mf.nc.gz','mks2.nc.gz','mm.nc.gz',
        'mn4.nc.gz','ms4.nc.gz','msf.nc.gz','msqm.nc.gz','mtm.nc.gz',
        'mu2.nc.gz','n2.nc.gz','n4.nc.gz','nu2.nc.gz','o1.nc.gz','p1.nc.gz',
        'q1.nc.gz','r2.nc.gz','s1.nc.gz','s2.nc.gz','s4.nc.gz','sa.nc.gz',
        'ssa.nc.gz','t2.nc.gz']
    c = ['2n2','eps2','j1','k1','k2','l2','lambda2','m2','m3','m4','m6',
        'm8','mf','mks2','mm','mn4','ms4','msf','msqm','mtm','mu2','n2',
        'n4','nu2','o1','p1','q1','r2','s1','s2','s4','sa','ssa','t2']
    model_format = 'FES'
    TYPE = 'z'
    SCALE = 1.0/100.0


# #### Setup coordinates for calculating tides

# In[ ]:


#-- create an image around the Ross Ice Shelf
xlimits = [-740000,520000]
ylimits = [-1430000,-300000]
spacing = [5e3,-5e3]
#-- 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 = np.int((xlimits[1]-xlimits[0])/spacing[0])+1
ny = np.int((ylimits[0]-ylimits[1])/spacing[1])+1
#-- convert image coordinates from polar stereographic to latitude/longitude
crs1 = pyproj.CRS.from_string("epsg:{0:d}".format(3031))
crs2 = pyproj.CRS.from_string("epsg:{0:d}".format(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 = datepick.value
tide_time = pyTMD.time.convert_calendar_dates(YMD.year, YMD.month,
    YMD.day, hour=np.arange(24))

#-- read tidal constants and interpolate to grid points
if model_format in ('OTIS','ATLAS'):
    amp,ph,D,c = extract_tidal_constants(lon, lat, grid_file, model_file,
        EPSG, TYPE=TYPE, METHOD='spline', GRID=model_format)
    DELTAT = np.zeros_like(tide_time)
elif (model_format == 'netcdf'):
    amp,ph,D,c = extract_netcdf_constants(lon, lat, model_directory,
        grid_file, model_files, TYPE=TYPE, METHOD='spline', SCALE=SCALE)
    DELTAT = np.zeros_like(tide_time)
elif (model_format == 'GOT'):
    amp,ph = extract_GOT_constants(lon, lat, model_directory, model_files,
        METHOD='spline', SCALE=SCALE)
    #-- interpolate delta times from calendar dates to tide time
    delta_file = pyTMD.utilities.get_data_path(['data','merged_deltat.data'])
    DELTAT = calc_delta_time(delta_file, tide_time)
elif (model_format == 'FES'):
    amp,ph = extract_FES_constants(lon, lat, model_directory, model_files,
        TYPE=TYPE, VERSION=MODEL, METHOD='spline', SCALE=SCALE)
    #-- interpolate delta times from calendar dates to tide time
    delta_file = pyTMD.utilities.get_data_path(['data','merged_deltat.data'])
    DELTAT = calc_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 map calculated every hour
tide_cm = np.ma.zeros((ny,nx,24))
for hour in range(24):
    #-- predict tidal elevations at time and infer minor corrections
    TIDE = predict_tide(tide_time[hour], hc, c, DELTAT=DELTAT[hour],
        CORRECTIONS=model_format)
    MINOR = infer_minor_corrections(tide_time[hour], hc, c,
        DELTAT=DELTAT[hour], CORRECTIONS=model_format)
    #-- add major and minor components and reform grid
    #-- convert from meters to centimeters
    tide_cm[:,:,hour] = 100.0*np.reshape((TIDE+MINOR),(ny,nx))


# #### Create animation of hourly tidal oscillation

# In[ ]:


#-- output Ross Ice Shelf Tide Animation
projection = ccrs.Stereographic(central_longitude=0.0,
    central_latitude=-90.0,true_scale_latitude=-71.0)
fig, ax = plt.subplots(num=1, figsize=(9,8),
    subplot_kw=dict(projection=projection))
#-- plot tide height
vmin,vmax = (np.min(tide_cm), np.max(tide_cm))
extent = (xlimits[0],xlimits[1],ylimits[0],ylimits[1])
im = 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')

#-- Add colorbar and adjust size
#-- pad = distance from main plot axis
#-- extend = add extension triangles to upper and lower bounds
#-- options: neither, both, min, max
#-- shrink = percent size of colorbar
#-- aspect = lengthXwidth aspect of colorbar
cbar = plt.colorbar(im, ax=ax, pad=0.025, extend='both',
    extendfrac=0.0375, shrink=0.85, aspect=22.5, drawedges=False)
#-- rasterized colorbar to remove lines
cbar.solids.set_rasterized(True)
#-- Add label to the colorbar
cbar.ax.set_ylabel('{0} Tide Height'.format(MODEL), fontsize=13)
cbar.ax.set_xlabel('cm', fontsize=13)
cbar.ax.xaxis.set_label_coords(0.50, 1.04)
#-- 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 = ax.set_title(None, fontsize=13)
#-- 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')
# adjust subplot within figure
fig.subplots_adjust(left=0.02,right=0.98,bottom=0.05,top=0.98)
           
#-- animate each map
def animate_maps(hour):
    #-- set map data
    im.set_data(tide_cm[:,:,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())


# In[ ]: