#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Thu Jun 11 16:29:29 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
from geoplots import mapplot
from amipHadISST.data_TCgenesis_map_mclim import mclim as tc
from amipHadISST.data_pTang_map_mclim import mclim as p
from amipHadISST.data_seedgenesis_map_mclim_rainy import mclim as seed
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
sel_region = lambda x: x.sel(lon=slice(260, 360), lat=slice(0,30)).where(x.lon+x.lat*35/20>295)
sel_region_r = lambda x: x.sel(lon=slice(260, 360), lat=slice(30,0)).where(x.lon+x.lat*35/20>295) #for ERA5 p

 
 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    import calendar
    #figname = __file__.replace('.py', f'_{tt.today()}.png')
    fig, axes = plt.subplots(3, 3, figsize=(8*1.1,5), sharex=True, sharey=True)
    months = range(6, 8+1)
    # row of TC genesis density
    da = tc.pipe(sel_region)
    levels=2.**np.arange(-4, 1.1, 0.5)
    cmap = 'OrRd'
    cbar_on = False
    label = 'TC gen. density'
    for i,mon in enumerate(months):
        ax = axes[0, i]
        im = da.sel(month=mon).plot.contourf(ax=ax, cmap=cmap, levels=levels, add_colorbar=cbar_on)
        plt.sca(ax)
        mapplot()
    fig.colorbar(im, ax=axes[0,:], label=label)

    #row of TC genesis probability
    da = p.pipe(sel_region)
    levels=np.arange(0, 0.51, 0.05)
    cmap = 'Spectral_r'
    cbar_on = False
    label = 'TC gen. prob.'
    for i,mon in enumerate(months):
        ax = axes[1, i]
        im = da.sel(month=mon).plot.contourf(ax=ax, cmap=cmap, levels=levels, add_colorbar=cbar_on)
        plt.sca(ax)
        mapplot()
    fig.colorbar(im, ax=axes[1,:], label=label)
    

    #row of seed genesis density
    da = seed.pipe(sel_region)
    levels=2.**np.arange(-2-1, 3.1-1, 0.5)
    cmap = 'OrRd'
    cbar_on = False
    label = 'seed gen. density'
    for i,mon in enumerate(months):
        ax = axes[2, i]
        im = da.sel(month=mon).plot.contourf(ax=ax, cmap=cmap, levels=levels, add_colorbar=cbar_on)
        plt.sca(ax)
        mapplot()
    fig.colorbar(im, ax=axes[2,:], label=label)


    for ax,s in zip(axes.flat, list('abcdefghi')):
        ax.set_title('')
        ax.set_xlabel('')
        ax.set_ylabel('')
        ax.text(0, 0, f'({s})', transform=ax.transAxes, va='bottom')
    for mon,ax in zip(months, axes[0,:]):
        ax.set_title(calendar.month_abbr[mon])

    if len(sys.argv)>1 and sys.argv[1]=='savefig':j
        figname = __file__.replace('.py', '.png')
        wysavefig(figname)
    tt.check(f'**Done**')