#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Sat May 22 17:05:46 EDT 2021
if __name__ == '__main__':
    from misc.timer import Timer
    tt = Timer(f'start {__file__}')
import sys, os.path, os, glob, datetime
import xarray as xr, numpy as np, pandas as pd, matplotlib.pyplot as plt
#more imports
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
yearsRef = slice(1800, 1849)
ifile = 'PMIP3_indices.nc'
ofile = ifile.replace('.nc', '_yearly.nc')
if os.path.exists(ofile):
    ds = xr.open_dataset(ofile)
    print('[loaded]:', ofile)
    #print('[exists]:', ofile)
    #sys.exit()
else:
#ssts
    ds = xr.open_dataset(ifile)
    models = [d.split('_')[0] for d in list(ds.dims)]#['CCSM4', 'CSIRO', 'GISS', 'HadCM3', 'IPSL', 'MIROC', 'MPI-ESM-P', 'bcc']
    models.sort(key=lambda x: x.upper())
    print(models)
    mdrs = []
    trops = []
    #monthly to yearly
    for model in models:
        da = ds[f'{model}_Atlantic'].groupby(f'{model}_Time.year').mean(f'{model}_Time')
        mdrs.append(da)
        da = ds[f'{model}_Tropical'].groupby(f'{model}_Time.year').mean(f'{model}_Time')
        trops.append(da)
    print('concat...')
    mdrs = xr.concat(mdrs, dim=pd.Index(models, name='model'))
    trops = xr.concat(trops, dim=pd.Index(models, name='model'))
    ds = xr.Dataset(dict(MDR=mdrs, TROP=trops))
    ds = ds.sel(year=slice(850,2000))# bcc stops at 2001-08, so ignore year 2001
    #sst anomalies
    if yearsRef is not None:
        MDRa = ds.MDR - ds.MDR.sel(year=yearsRef).mean('year')
        TROPa = ds.TROP - ds.TROP.sel(year=yearsRef).mean('year')
    else:
        MDRa = ds.MDR - ds.MDR.mean('year')
        TROPa = ds.TROP - ds.TROP.mean('year')
    #HU (hurricane #)
    HU = np.exp(1.707 + 1.388*MDRa - 1.521*TROPa)
    HU.attrs['long_name'] = 'HU#: EXP(1.707 + 1.388*MDRa - 1.521*TROPa)'
    if yearsRef:
        HU.attrs['yearsRef'] = f'{yearsRef.start}-{yearsRef.stop}'
    ds['HU'] = HU
    #TS (tropical storm #)
    TS = np.exp(2.10356 + 0.97612*MDRa - 0.97102*TROPa)
    TS.attrs['long_name'] = 'TS#: EXP(2.10356 + 0.97612*MDRa - 0.97102*TROPa)'
    if yearsRef:
        TS.attrs['yearsRef'] = f'{yearsRef.start}-{yearsRef.stop}'
    ds['TS'] = TS
    #PDI
    PDI = np.exp(0.76 + 1.87*MDRa - 1.58*TROPa)
    PDI.attrs['long_name'] = 'PDI: EXP(0.76 + 1.87*MDRa - 1.58*TROPa)'
    if yearsRef:
        PDI.attrs['yearsRef'] = f'{yearsRef.start}-{yearsRef.stop}'
    ds['PDI'] = PDI
    print('saving...')
    ds.to_netcdf(ofile)
    print('[saved]:', ofile)
 
 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    import xfilter
    da = ds.HU.interpolate_na(dim='year', method='nearest', fill_value='extrapolate') \
        .filter.lowpass(1/40, dim='year', padtype='constant') \
        .where(ds.HU*0==0)
    zz_sample = np.random.poisson(lam=ds.HU.fillna(0).values, size=(10000,)+ds.HU.shape)
    da_sample = xr.DataArray(zz_sample, dims=('sample',)+da.dims).assign_coords(model=da.model, year=da.year) \
        .where(ds.HU*0==0) \
        .interpolate_na(dim='year', method='nearest', fill_value='extrapolate') \
        .filter.lowpass(1/40, dim='year', padtype='constant') \
        .where(ds.HU*0==0)
    q025 = da_sample.quantile(0.025, dim='sample').drop('quantile')
    q975 = da_sample.quantile(0.975, dim='sample').drop('quantile')
    #mystd = da_sample.std('sample')
    g = da.plot(col='model', col_wrap=2, figsize=(11,6), xlim=(850,2000), xticks=range(900,2001,100))
    for ax,model in zip(g.axes.flat, ds.model.values):
        ax.fill_between(da.year, q025.sel(model=model), q975.sel(model=model), alpha=0.3)
        #ax.fill_between(da.year, da.sel(model=model) - mystd.sel(model=model), da.sel(model=model) + mystd.sel(model=model), alpha=0.3, color='C1')
        ax.grid('on')
    
    #savefig
    if 'savefig' in sys.argv:
        figname = __file__.replace('.py', f'.png')
        wysavefig(figname)
    tt.check(f'**Done**')
    plt.show()