#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Wed Feb 17 10:39:54 EST 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
import matplotlib.pyplot as plt
#more imports
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
ifile = 'HadISST.sstindex.amean.187001-201912.nc'
ds = xr.open_dataset(ifile)
def sst2hu(ds, mdr, trop, yearsRef):
    ssta_mdr = ds[mdr] - ds[mdr].sel(year=yearsRef).mean()
    ssta_trop = ds[trop] - ds[trop].sel(year=yearsRef).mean()
    hu = np.exp(1.707 + 1.388*ssta_mdr - 1.521*ssta_trop)
    hu.attrs['long_name'] = 'HU#: EXP(1.707 + 1.388*ssta_mdr - 1.521*ssta_trop)'
    hu.attrs['yearsRef'] = f'{yearsRef.start}-{yearsRef.stop}'

    return hu
    
 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    import xfilter
    def lowpass(da): return da.filter.lowpass(1/40, dim='year', padtype='even')
    fig, ax = plt.subplots(figsize=(8,4))
    lww = 2
    lw = 1

    yearsRef = slice(1982,2005)
    hu = sst2hu(ds, 'sst_mdr', 'sst_trop', yearsRef)
    hu.plot(ax=ax, label=hu.attrs['yearsRef'], color='k', lw=lww)
    hu.pipe(lowpass).plot(ax=ax, color='k', lw=lw)

    yearsRef = slice(1901,2000)
    hu = sst2hu(ds, 'sst_mdr', 'sst_trop', yearsRef)
    hu.plot(ax=ax, label=hu.attrs['yearsRef'], color='gray', lw=lww)
    hu.pipe(lowpass).plot(ax=ax, color='gray', lw=lw)

    yearsRef = slice(1979,2000)
    hu = sst2hu(ds, 'sst_mdr', 'sst_trop', yearsRef)
    hu.plot(ax=ax, label=hu.attrs['yearsRef'], color='C0', ls='--', lw=lww)
    hu.pipe(lowpass).plot(ax=ax, color='C0', ls='--', lw=lw, label='40-year lowpass')

    ax.legend(title='reference years', ncol=2, frameon=False)
    ax.set_ylabel('HadISST emulated NA HU #')
    
    figname = __file__.replace('.py', f'_{tt.today()}.png')
    if len(sys.argv) > 1 and sys.argv[1] == 'savefig':
        plt.savefig(figname)
        print('[saved]:', figname)
    tt.check(f'**Done**')
    plt.show()