#!/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 = 'density_g'
long_name  = 'ERA5 storm12hOcean genesis density'
units = '# per month per 10x10 box'
years = slice('1981', '2018')
figsize = (8,5)
bs = tc_basins()
ifile = '/tigress/wenchang/analysis/TC/ERA5/exp_p75C_550k_48h_12ms/netcdf/ERA5.tc_density.allstorms.12hOcean.1979-2018.nc'
if 'da' not in locals():
    ds = xr.open_dataset(ifile)
    da = ds[data_name].sel(time=years).groupby('time.month').mean('time')

plt.close('all')
g = da.where(bs.mask(da)==bs.map_keys(basin)) \
    .sel(lat=slice(10, 30), lon=slice(260, 360)) \
    .assign_attrs(long_name=long_name, units=units) \
    .plot(col='month', col_wrap=3, cmap='Spectral_r', levels=21, vmax=2, figsize=figsize)
for ax in g.axes.flatten():
    plt.sca(ax)
    #mapplot(coastlines_color='lightgray')#, coastlines_width=0.5, xticks=range(0,360,120))
    #mapplot(fill_continents=True, continents_color='w')#, coastlines_width=0.5, xticks=range(0,360,120))
    mapplot(coastlines_color='lightgray', coastlines_width=0.5)
    ax.set_xlabel('')
    ax.set_ylabel('')
plt.tight_layout(rect=(0,0,0.84,1))

figname = f'fig.map.genesis_density.storm12hOcean.mclim.ERA5.{years.start}-{years.stop}.png'
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()