#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Thu Sep  9 15:17:30 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
ds_obs = xr.open_dataset('humh_adjusted.nc')
ds_hadisst = xr.open_dataset('HadISST_ssttc.1870-2019.nc')
ds_chan = xr.open_dataset('HadISST_ssttc_plusChanCorr.1870-2019.nc')
if len(sys.argv)>1:
    try:
        lpwindow = int(sys.argv[1]) #lpwindow is specified from the first arg
    except:
        lpwindow = None
else:
    lpwindow = None#40#15#5#None 
if lpwindow is not None:
    lowpass = lambda x: x.filter.lowpass(1/lpwindow, dim='year', padtype='even')
else:
    lowpass = lambda x: x
 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    alpha = 0.1
    qq = [0.025, 0.975]
    plt.close()
    fig, ax = plt.subplots()
    #raw obs
    da = ds_obs.HU_raw.pipe(lowpass)
    da.plot(ls='--', color='gray', label='raw HU records')
    #corrected obs
    #ci
    da = ds_obs.HU.pipe(lowpass)
    ci = da.quantile(qq, dim='sample')
    ax.fill_between(ci.year, ci[0], ci[1], color='k', alpha=alpha)
    #mean
    da = ds_obs.HU.pipe(lowpass).mean('sample')
    da.plot(color='k', label='corrected HU records')
    
    #sst model: hadisst
    da = ds_hadisst.HU
    rng = np.random.default_rng(0)
    zz = rng.poisson(da, size=(10000,)+da.shape)
    da_ = xr.DataArray(zz, dims=('sample',)+da.dims).assign_coords(year=da.year.values)
    ci = da_.pipe(lowpass).quantile(qq, dim='sample')
    ax.fill_between(ci.year, ci[0], ci[1], color='C0', alpha=alpha)
    da.pipe(lowpass).plot(color='C0', label='HadISST emulation')
    #sst model: corrected hadisst
    da = ds_chan.HU
    #rng = np.random.default_rng(0)
    zz = rng.poisson(da, size=(10000,)+da.shape)
    da_ = xr.DataArray(zz, dims=('sample',)+da.dims).assign_coords(year=da.year.values)
    ci = da_.pipe(lowpass).quantile(qq, dim='sample')
    ax.fill_between(ci.year, ci[0], ci[1], color='C1', alpha=alpha)
    da.pipe(lowpass).plot(color='C1', label='corrected HadISST emulation')

    ax.legend(ncol=2, loc='upper left')
    title = 'HU annual count'
    if lpwindow is not None:
        title = title + f' ({lpwindow}-year lowpass)' 
    ax.set_title(title, loc='left')
    ax.set_ylabel('HU #')
    ax.set_xlim(1850, 2020)
    ax.set_xticks(range(1860,2021,20))
    if lpwindow is None:
        ax.set_ylim(0,24)
        ax.set_yticks(range(0,25,4))
    elif lpwindow == 5:
        ax.set_ylim(2,16)
        ax.set_yticks(range(2,17,2))
    elif lpwindow == 15:
        ax.set_ylim(3,13)
        ax.set_yticks(range(3,14,1))
    elif lpwindow == 40:
        ax.set_ylim(3.5,12)
        ax.set_yticks(range(4,13))
    
    #savefig
    if 'savefig' in sys.argv:
        if lpwindow is None:
            figname = __file__.replace('.py', f'.png')
        else:
            figname = __file__.replace('.py', f'_lp{lpwindow}year.png')
        wysavefig(figname)
    tt.check(f'**Done**')
    plt.show()