#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Wed Jun 10 14:52:28 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 xtc import tc_density
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
savedfile = os.path.abspath(__file__).replace('.py', '.nc')
dataname = 'p'
long_name = 'TC genesis probability'

if os.path.exists(savedfile): # read from saved file
    print('**reading from saved file**:', savedfile)
    with xr.open_dataset(savedfile) as ds:
        da = ds[dataname]
else: # calculate from raw tracks data
    ifile = os.path.join(os.path.dirname(__file__), 'data.VI.nc')
    with xr.open_dataset(ifile) as ds:
        vi = ds['VI']
        vi0 = 0.014
        n = np.log(9)/np.log(0.1/0.014)
        da = 1./(1. + (vi/vi0)**n)
        da = da.assign_attrs(long_name=long_name)
        da.to_dataset(name=dataname).to_netcdf(savedfile)
        print('**saved**:', savedfile)

 
 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    from geoplots import mapplot
    fig, ax = plt.subplots()
    years = slice('1980', '2018')
    #figname = __file__.replace('.py', f'.aclim{years.start}-{years.stop}.{tt.today()}.png')
    da.sel(time=years).groupby('time.year').mean('time').mean('year').mean('en') \
        .assign_attrs(long_name=da.attrs['long_name'])\
        .plot.contourf(cmap='Spectral_r', levels=np.arange(0, 0.41, 0.025), cbar_kwargs=dict(aspect=30))
    mapplot()
    ax.set_title(f'AM2.5 {years.start}-{years.stop} TC genesis probability')
    #ax.axhline(10, color='gray', ls='--')
    #ax.axhline(-10, color='gray', ls='--')

    if len(sys.argv)>1 and sys.argv[1]=='savefig':
        figname = __file__.replace('.py', f'.aclim{years.start}-{years.stop}.png')
        wysavefig(figname)

    tt.check(f'**Done**')
    plt.show()