#!/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
from misc.cim import cim, sem
if __name__ == '__main__':
    tt.check('end import')

# MDR lons and lats
# ventilation index
def get_pTang(basin='NA'):
    '''ventilation index to TC genesis probability (based on Tang and Emanuel'''
    thisdir = os.path.dirname(__file__)
    ofile = os.path.join(thisdir, f'data_pTang_{basin}.nc')
    if os.path.exists(ofile):
        print('[exists]:', ofile)
        return xr.open_dataset(ofile)['p']
    ifile = os.path.join(thisdir, 'data.VI.nc')
    bs = tc_basins()
    vi = xr.open_dataset(ifile)['VI']
    vi = (vi
        # NA basin
        .where( bs.mask(vi)==bs.map_keys(basin) )
        # lat range 10-30N
        .where(vi.lat>10).where(vi.lat<30)
        )
    vi0 = 0.014
    n = np.log(9)/np.log(0.1/0.014)
    p = vi.pipe(lambda x: 1./(1. + (x/vi0)**n)).geo.fldmean()
    long_name = 'TC genesis probability'
    p = p.assign_attrs(long_name=long_name)
    p.to_dataset(name='p').to_netcdf(ofile)
    print('[saved]:', ofile)

    return p 

#annual cycle
years = slice('1980', '2018')
def get_cycle(basin='NA', years=years):
    da = get_pTang(basin=basin)
    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 = 'EP'
    model = 'HIRAM'
    expname = 'amipHadISST'
    figname = f'data_pTang_{basin}_{expname}_{tt.today()}.png'
    da = get_pTang(basin=basin)
    da.groupby('time.year').mean('time').plot(hue='en')
    """
    figname = f'data_pTang_cycle_{basin}_{expname}_{tt.today()}.png'
    ds_ = get_cycle(basin=basin, years=years)
    long_name = f'pTang({basin}, {expname})'
    ax = ds_.rename(mclim=long_name).to_dataframe().plot.bar(y=long_name, yerr='sem', rot=0)
    ax.legend(loc='upper left')
    """
    plt.savefig(figname)
    print('[saved]:', figname)

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