#!/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', fixterm=None):
    '''ventilation index to TC genesis probability (based on Tang and Emanuel.
    fixterm: whether Vshear, chi or PI is replaced with annual mean climatology.'''
    thisdir = os.path.abspath( os.path.dirname(__file__) )
    if fixterm is None:
        ofile = f'{thisdir}/data_pTang_{basin}.nc'
    else:
        ofile = f'{thisdir}/data_pTang_const{fixterm}_{basin}.nc'
    if os.path.exists(ofile):
        print('[exists]:', ofile)
        return xr.open_dataset(ofile)['p']
    bs = tc_basins()
    if fixterm == 'Vshear':
        Vshear = xr.open_dataset(f'{thisdir}/data.Vshear.nc')['Vshear'].sel(time=slice('1980','2018')).mean('time').mean('en')
        chi = xr.open_dataset(f'{thisdir}/data.chi.nc')['chi']#.sel(time=slice('1980','2018')).mean('time').mean('en')
        V_PI = xr.open_dataset(f'{thisdir}/data.PI.nc')['vmax']#.sel(time=slice('1980','2018')).mean('time').mean('en')
        vi = Vshear * chi /V_PI
    elif fixterm == 'chi':
        Vshear = xr.open_dataset(f'{thisdir}/data.Vshear.nc')['Vshear']#.sel(time=slice('1980','2018')).mean('time').mean('en')
        chi = xr.open_dataset(f'{thisdir}/data.chi.nc')['chi'].sel(time=slice('1980','2018')).mean('time').mean('en')
        V_PI = xr.open_dataset(f'{thisdir}/data.PI.nc')['vmax']#.sel(time=slice('1980','2018')).mean('time').mean('en')
        vi = Vshear * chi /V_PI
    elif fixterm == 'PI':
        Vshear = xr.open_dataset(f'{thisdir}/data.Vshear.nc')['Vshear']#.sel(time=slice('1980','2018')).mean('time').mean('en')
        chi = xr.open_dataset(f'{thisdir}/data.chi.nc')['chi']#.sel(time=slice('1980','2018')).mean('time').mean('en')
        V_PI = xr.open_dataset(f'{thisdir}/data.PI.nc')['vmax'].sel(time=slice('1980','2018')).mean('time').mean('en')
        vi = Vshear * chi /V_PI
    elif fixterm == 'chi+PI':
        Vshear = xr.open_dataset(f'{thisdir}/data.Vshear.nc')['Vshear']#.sel(time=slice('1980','2018')).mean('time').mean('en')
        chi = xr.open_dataset(f'{thisdir}/data.chi.nc')['chi'].sel(time=slice('1980','2018')).mean('time').mean('en')
        V_PI = xr.open_dataset(f'{thisdir}/data.PI.nc')['vmax'].sel(time=slice('1980','2018')).mean('time').mean('en')
        vi = Vshear * chi /V_PI
    else:
        ifile = f'{thisdir}/data.VI.nc'
        vi = xr.open_dataset(ifile)['VI']
    vi.load()
    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 


if __name__ == '__main__':
    from wyconfig import * #my plot settings
    basin = 'NA'
    model = 'AM2.5'
    expname = 'amipHadISST'
    figname = __file__.replace('.py', f'_{basin}_{tt.today()}.png')
    dc = {v:get_pTang(basin=basin, fixterm=v) for v in [None, 'Vshear', 'chi', 'PI', 'chi+PI']}
    for key,da in dc.items():
        if key is None:
            label = 'ctl'
            color = 'k'
        elif key == 'chi+PI':
            label = f'const $\chi +$ PI'
            color = None
        else:
            label = f'const {key}'
            color = None
        da.groupby('time.year').mean('time').mean('en').plot(label=label, color=color)
    plt.legend()
    title = f'{model} {expname} {basin} TC genesis probability'
    plt.title(title)


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

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