#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Thu May 20 17:12:35 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
import xfilter
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
ifile = 'hu.0850-2100ssp585.MRI-ESM2-0.nc'
ds = xr.open_dataset(ifile).load()
model = ifile.split('.')[-2]#MRI-ESM2-0
da = ds.HU
#spread of the lowpass time series
zz = da.values
zz_pn = []
for lam in zz:
    zz_pn.append(np.random.poisson(lam, size=10000))
da_pn = xr.DataArray(zz_pn, dims=['year', 'sample']).assign_coords(year=da.year.values)
std = da_pn.filter.lowpass(1/40, dim='year', padtype='even').std('sample')
q975 = da_pn.filter.lowpass(1/40, dim='year', padtype='even').quantile(0.975, dim='sample').drop('quantile')
q025 = da_pn.filter.lowpass(1/40, dim='year', padtype='even').quantile(0.025, dim='sample').drop('quantile')
da = ds.HU.filter.lowpass(1/40, dim='year', padtype='even')
stde40 = da.pipe(lambda x: x**0.5/np.sqrt(40))
ds = xr.Dataset(dict(meanvalue=da, stdvalue=std, q975=q975, q025=q025, stde40=stde40))
 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    fig,ax = plt.subplots(figsize=(8,3))
    ax.fill_between(da.year, ds.q025, ds.q975, alpha=0.2, color='C0', label='quantile [0.025, 0.975]')
    ax.fill_between(ds.year, ds.meanvalue-ds.stdvalue, ds.meanvalue+ds.stdvalue, alpha=0.5, color='C0', label='$\pm$1 STD')
    ds.meanvalue.plot(ax=ax, label='mean(=$\lambda$)')
    (ds.meanvalue + ds.stde40).plot(ax=ax, label='mean + $\sqrt{\lambda}$/$\sqrt{40}$', ls='--', color='C0')
    ax.set_xticks(range(800,2101,100))
    ax.set_xlim(850,2100)
    ax.set_title(f'{model}: past1000 + historical + ssp585')
    ax.grid('on')
    ax.legend()

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