#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Wed Mar  4 19:10:44 EST 2020
if __name__ == '__main__':
    from misc.timer import Timer
    tt = Timer(f'start {__file__}')
import sys, os.path, os, datetime
import xarray as xr, numpy as np, pandas as pd
#import matplotlib.pyplot as plt
import geoxarray
#from xtc import tc_basins
maindir = '/tigress/wenchang/analysis/seedTC'
if maindir not in sys.path:
    sys.path.append(maindir)
from AM2p5C360.data_pTang_partial import get_pTang
from misc.cim import cim, sem
if __name__ == '__main__':
    tt.check('end import')


#annual cycle
def get_cycle(basin='NA', years=None, fixterm=None):
    if years is None:
        years = slice('1980', '2018')
    da = get_pTang(basin=basin, fixterm=fixterm)
    da = da.sel(time=years)
    # stack dims of ['en', 'time'] and use the stacked 'time' as the new coord
    if 'en' in da.dims:
        da = da.stack(s=[...]).pipe(lambda x: x.assign_coords(s=x.time)).rename(s='time')
    da_mclim = da.groupby('time.month').mean('time')
    da_cim = da.groupby('time.month').map(cim, dim='time')
    da_sem = da.groupby('time.month').map(sem, dim='time')
    ds_mclim = xr.Dataset(dict(
        mclim=da_mclim,
        cim=da_cim,
        sem=da_sem
        ))
    return ds_mclim 


if __name__ == '__main__':
    from wyconfig import * #my plot settings
    basin = 'NA'
    model = 'AM2.5C360'
    expname = 'amipHadISST'
    years = slice('1980', '2018')
    figname = __file__.replace('.py', f'_{basin}_{tt.today()}.png')
    dc = {v:get_cycle(basin=basin, years=years, fixterm=v) for v in [None, 'Vshear', 'chi', 'PI', 'chi+PI']}
    for key,ds in dc.items():
        if key is None:
            label = 'CTL'
            color = 'k'
        elif key == 'chi+PI':
            label = 'const $\chi +$ PI'
            color = None 
        else:
            label = f'const {key}'
            color = None
        ds.mclim.plot(label=label, color=color, marker='o')
    plt.legend(loc='upper left')
    plt.title(f'{model} {expname} {basin} TC genesis probability annual cycle')
    plt.xticks(range(1,13))
    plt.ylabel(f'p($\Lambda$)')

    plt.savefig(figname)
    print('[saved]:', figname)

    tt.check('**done**')
    plt.show()