#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Thu Jun  3 12:16:09 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
lowpass = lambda x: x.filter.lowpass(1/40, dim='year', padtype='even')
datafile = __file__.replace('.py', '.nc')
if os.path.exists(datafile):
    ds = xr.open_dataset(datafile)
    print('[loaded]:', datafile)
else:
    ifile = 'LMR2019_TC.nc'
    ds = xr.open_dataset(ifile).rename(MCrun='mcrun').sel(year=slice(850,None))
    da = ds.HU
    damean = da.pipe(lowpass).mean('mcrun')
    #uncertainty from sst
    print('sst...')
    stde = da.pipe(lowpass).std('mcrun')
    ci_sst = (damean-stde, damean+stde)
    stde_sst = stde
    #qq = [0.025, 0.975]
    rng = np.random.default_rng(0)
    #uncertainty from Poisson stochastics
    print('Poisson stochastics...')
    da_ = da.pipe(np.log).mean('mcrun').pipe(np.exp) #geometric mean
    zz = rng.poisson(da_, size=(10000,)+da_.shape)
    da_smp = xr.DataArray(zz, dims=('smp',)+da_.dims).assign_coords(year=da_.year.values)
    stde = da_smp.pipe(lowpass).std('smp')
    damean_gm = da_.pipe(lowpass)#geometric mean
    ci_ps = (damean_gm-stde, damean+stde)
    stde_ps = stde
    #uncertainty from both factors
    print('both...')
    zz = rng.poisson(da, size=(10000,)+da.shape)
    da_smp = xr.DataArray(zz, dims=('smp',)+da.dims).assign_coords(year=da.year.values, mcrun=da.mcrun.values)
    stde = da_smp.pipe(lowpass).stack(s=['smp','mcrun']).std('s')
    ci_both = (damean-stde, damean+stde)
    stde_both = stde

    ds = xr.Dataset(dict(meanHU=damean, gmeanHU=damean_gm, stde_both=stde_both, stde_sst=stde_sst, stde_ps=stde_ps))
    print('saving...')
    ds.to_netcdf(datafile)
    print('[saved]:', datafile)
 
damean = ds.meanHU
damean_gm = ds.gmeanHU
ci_both = damean - ds.stde_both, damean + ds.stde_both
ci_sst = damean - ds.stde_sst, damean + ds.stde_sst
ci_ps = damean_gm - ds.stde_ps, damean_gm + ds.stde_ps

if __name__ == '__main__':
    from wyconfig import * #my plot settings
    plt.close('all')
    #fig,ax = plt.subplots()
    fig,axes = plt.subplots(2, 1, figsize=(6,5))
    alpha = 0.3

    ax = axes[0]
    ax.fill_between(da.year, ci_both[0], ci_both[1], alpha=alpha, label='both')
    ax.fill_between(da.year, ci_sst[0], ci_sst[1], alpha=alpha, label='sst')
    ax.fill_between(da.year, ci_ps[0], ci_ps[1], alpha=alpha, label='Poisson')
    damean.plot(ax=ax)
    damean_gm.plot(color='C2', ls='--', ax=ax)
    ax.legend(ncol=3)
    ax.set_xlabel('')
    ax.set_ylabel('HU #')

    #fig,ax = plt.subplots()
    ax = axes[1]
    stde_both.plot(ax=ax, label='both')
    stde_sst.plot(ax=ax, label='sst')
    stde_ps.plot(ax=ax, label='Poisson stochastics')
    ( ( stde_ps**2 + stde_sst**2 )**0.5 ).plot(ls='--', ax=ax, label='sst + Poisson stochastics')
    ax.legend(ncol=2)
    ax.set_ylabel('HU # stde')

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