#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Wed Mar 11 16:43:12 EDT 2020
import sys, os.path, os, datetime
import xarray as xr, numpy as np, pandas as pd
import matplotlib.pyplot as plt
#more imports
from xtc.basins import tc_basins
from geoplots import mapplot
#
if __name__ == '__main__':
    print()
    today = datetime.date.today()
    today_s = today.strftime('%Y-%m-%d')
    tformat = '%Y-%m-%dT%H:%M:%S'
    t_start = datetime.datetime.now()
    print('[start]:', t_start.strftime(tformat))
    
#start from here
basin = 'NA'
data_name = 'VI'
long_name  = 'ERA5 pTK(VI)'
years = slice('1981', '2018')
bs = tc_basins()
ifile = 'data.VI.ERA5.nc'
if 'da' not in locals():
    ds = xr.open_dataset(ifile)
    n = np.log(9)/np.log(0.1/0.014)
    da = ds[data_name].sel(time=years).pipe(lambda x: 1./(1.+(x/0.014)**n)).groupby('time.month').mean('time')

plt.close('all')
g = da.where(bs.mask(da)==bs.map_keys(basin)) \
    .sel(lat=slice(30, 10), lon=slice(260, 360)) \
    .assign_attrs(long_name=long_name) \
    .plot(col='month', col_wrap=3, cmap='Spectral_r', levels=21, vmin=0, vmax=1, figsize=(8,5))
for ax in g.axes.flatten():
    plt.sca(ax)
    mapplot(coastlines_color='gray')
    ax.set_xlabel('')
    ax.set_ylabel('')
plt.tight_layout(rect=(0,0,0.84,1))

figname = 'fig.map.pTK.clim.NA.ERA5.png'
plt.show()
plt.savefig(figname)
 
if __name__ == '__main__':
    t_end = datetime.datetime.now()
    print('[end]:', t_end.strftime(tformat))
    print('[total time used]:', f'{(t_end - t_start).seconds:,} seconds')
    print()