#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Tue Jul 20 18:05:44 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
ifile = 'data/nino34.ctl.nyears1000.nc'
da = xr.open_dataarray(ifile)
da = da.assign_coords(time=xr.cftime_range('0000-01', freq='MS', periods=da.size, calendar='noleap'))
da_ = da.isel(time=slice(0, 300*12)).pipe(lambda x: x.groupby('time.month') - x.groupby('time.month').mean('time')) #select the first 300 years and re-normalize the time series
das = [da_.isel(time=slice(ii,ii+15*12)).drop(['time', 'month']) for ii in range(0, 286*12, 12)]
ts = xr.concat(das, dim=pd.Index(range(len(das)), name='ens')).assign_coords(time=pd.Index(pd.date_range('2000-01', freq='MS', periods=15*12)))
#MC resampling
rng = np.random.default_rng(0)
ts_mc = xr.DataArray(rng.choice(ts, size=(30,1000)).mean(axis=0), dims=['mc', 'time']).assign_coords(time=ts.time.values)
q = ts_mc.quantile([.025, .975], dim='mc')

 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    fig, ax = plt.subplots()
    ax.fill_between(q.time.values, q.isel(quantile=0), q.isel(quantile=1), alpha=.25)
    ts_mc.mean('mc').plot()

    times = ts_mc.time.values[0::12]
    ax.set_xticks(times)
    ax.set_xticklabels([f'{ii:02d}' for ii in range(15)])
    ax.set_ylabel('Nino3.4 index [$^\circ$C]')
    ax.set_xlabel('year')
    ax.set_ylim(-1.5,1.5)

    
    #savefig
    if 'savefig' in sys.argv:
        figname = __file__.replace('.py', f'.png')
        wysavefig(figname)
    tt.check(f'**Done**')
    plt.show()