#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Thu Sep  9 17:30:57 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
if 'TROP' in sys.argv:
    daname = 'TROP'
else:
    daname = 'MDR'
ds_obs = xr.open_dataset('HadISST_sstmdrtrop_plusChanCorr.1870-2019.nc')
ds_lmr2018 = xr.open_dataset('LMR2018_TC.nc')
ds_lmr2019 = xr.open_dataset('LMR2019_TC.nc')
anomaly = lambda x: x - x.sel(year=slice(1982,2005)).mean()
lpwindow = 40
lowpass = lambda x: x.filter.lowpass(1/lpwindow, dim='year', padtype='even')
selyears = lambda x: x.sel(year=slice(1870, None))
def normalize(da):
    yearsRef = slice(1901,2000)
    da_lp = da.pipe(lowpass).mean('MCrun')
    selfMean = da_lp.sel(year=yearsRef).mean('year')
    selfSTD = da_lp.sel(year=yearsRef).std('year')
    da_target = ds_obs[daname].pipe(anomaly).pipe(lowpass) 
    targetMean = da_target.sel(year=yearsRef).mean('year')
    targetSTD = da_target.sel(year=yearsRef).std('year')
    return (da - selfMean)*targetSTD/selfSTD + targetMean
 
 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    #plt.close()
    alpha = 0.1
    fig, ax = plt.subplots()
    #lmr2018
    da = ds_lmr2018[daname]
    da.pipe(anomaly).pipe(lowpass).pipe(selyears).plot(hue='MCrun', color='C0', alpha=alpha, add_legend=False, ax=ax)
    da.pipe(anomaly).pipe(lowpass).pipe(selyears).mean('MCrun').plot(color='C0', ax=ax, label='LMR2.0')
    #lmr2019
    da = ds_lmr2019[daname]
    da.pipe(anomaly).pipe(lowpass).pipe(selyears).plot(hue='MCrun', color='C1', alpha=alpha, add_legend=False, ax=ax)
    da.pipe(anomaly).pipe(lowpass).pipe(selyears).mean('MCrun').plot(color='C1', ax=ax, label='LMR2.1')
    #obs
    da = ds_obs[daname]
    da.pipe(anomaly).pipe(lowpass).plot(color='k', ax=ax, label='HadISSTchan')

    #lmr2018 transformed
    da = ds_lmr2018[daname]
    da.pipe(normalize).pipe(lowpass).pipe(selyears) \
        .plot(hue='MCrun', color='C2', alpha=alpha, add_legend=False, ax=ax)
    da.pipe(normalize).pipe(lowpass).pipe(selyears) \
        .mean('MCrun') \
        .plot(color='C2', ls='--', ax=ax, label='LMR2.0 transformed')
    #lm2018 transformed
    da = ds_lmr2019[daname]
    da.pipe(normalize).pipe(lowpass).pipe(selyears) \
        .plot(hue='MCrun', color='C3', alpha=alpha, add_legend=False, ax=ax)
    da.pipe(normalize).pipe(lowpass).pipe(selyears) \
        .mean('MCrun') \
        .plot(color='C3', ls='--', ax=ax, label='LMR2.1 transformed')


    ax.legend()
    ax.set_ylabel(f'{daname} [K]')
    ax.set_ylim(-0.6, 0.4)
    ax.set_xlim(1870, 2020)
    
    #savefig
    if 'savefig' in sys.argv:
        figname = __file__.replace('.py', f'_{daname}.png')
        wysavefig(figname)
    tt.check(f'**Done**')
    plt.show()