#!/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__)))
if maindir not in sys.path:
    sys.path.append(maindir)
from ERA5.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
    basin = 'NA'
    lats, lons = slice(0, 40), slice(260, 360)# NA basin
    figname = __file__.replace('.py', f'.NA.{tt.today()}.png')
    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).sel(lon=lons, lat=lats)\
            .pipe(lambda x: x.where(x.lon+x.lat*35/20>295)) \
            .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.98, month_abbr[i+1], va='top', transform=ax.transAxes)

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

    plt.savefig(figname)
    print('**saved**:', figname)
    tt.check(f'**Done**')
    plt.show()