#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Wed Jun 10 15:39:05 EDT 2020
if __name__ == '__main__':
    from misc.timer import Timer
    tt = Timer(f'start {__file__}')
import sys, os.path, os, glob
import xarray as xr, numpy as np, pandas as pd
#import matplotlib.pyplot as plt
#more imports
maindir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
maindir = os.path.dirname(maindir)
if maindir not in sys.path:
    sys.path.append(maindir)
from IBTrACS.data_TCgenesis_map import da as gdensity
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
years = slice('1980', '2018')
da = gdensity
mclim = da.sel(time=years).groupby('time.month').mean(['time', ])
mclim = mclim.assign_attrs(long_name=da.attrs['long_name'], units=da.attrs['units'])
 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    from calendar import month_abbr
    from geoplots import mapplot
    basin = 'SI'
    lats, lons = slice(-40, 0), slice(30, 105)# SI basin
    fig, axes = plt.subplots(4, 3, figsize=(8, 5), sharex=True, sharey=True)
    for i,ax in enumerate(axes.flat):
        im = mclim.isel(month=i-6).sel(lon=lons, lat=lats)\
            .plot.contourf(ax=ax, cmap='OrRd', levels=2.**np.arange(-4, 1.1, 0.5), add_colorbar=False)
        plt.sca(ax)
        mapplot()
        ax.set_title('')
        ax.set_xlabel('')
        ax.set_ylabel('')
        ax.text(0.02, 0.02, month_abbr[i+1-6 if i+1-6>0 else i+1-6-1], va='bottom', transform=ax.transAxes)

    fig.colorbar(im, ax=axes, aspect=30,
        label=f"IBTrACS {years.start}-{years.stop} {mclim.attrs['long_name']}\n[{mclim.attrs['units']}]")
    #fig.suptitle(f"{mclim.attrs['long_name']}: {years.start}-{years.stop} mclim")

    if len(sys.argv)>1 and 'savefig' in sys.argv[1:]:
        figname = __file__.replace('.py', f'.{basin}.{tt.today()}.png')
        plt.savefig(figname)
        print('**saved**:', figname)
    tt.check(f'**Done**')
    plt.show()