#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Mon Jul 21 03:06:49 PM EDT 2025
if __name__ == '__main__':
    import sys,os
    try:
        from misc.timer import Timer
        tt = Timer(f'[{os.getcwd()}] start ' + ' '.join(sys.argv))
    except:
        pass
import sys, os.path, os, glob, datetime
import xarray as xr, numpy as np, pandas as pd, matplotlib.pyplot as plt
#more imports
wython = '/tigress/wenchang/wython'
if wython not in sys.path: sys.path.append(wython); print('added to python path:', wython)
from misc import get_kws_from_argv
import xfilter
#
if __name__ == '__main__':
    try:
        tt.check('end import')
    except:
        pass
#
#start from here
basin = 'NA'
timeslice_ref = slice('1351', '1550')
yearslice_ref = slice(1351, 1550)
norm = lambda da: da/da.sel(year=yearslice_ref).mean('year')*100 - 100
norm_bound = lambda da: da/da_lizzie.sel(year=yearslice_ref).mean('year')*100 - 100
anom = lambda da: da - da.sel(year=yearslice_ref).mean('year')
units = lambda da: da.assign_attrs(units=f'% of {timeslice_ref.start}-{timeslice_ref.stop} mean', long_name='HU')
pct_diff = lambda da: da.sel(year=slice(1451, 1550)).mean('year') - da.sel(year=slice(1351,1450)).mean('year')
#lowpass
n_window = 40 #40 years
lowpass = lambda da: da.filter.lowpass(1/n_window, dim='year', padtype='even')

#sediment HU
#ifile = '/tigress/wenchang/analysis/LMHU/data/Lizzie/sedimentHU_v20220711ens_noCaySal_wy_max_lp40_count.nc'
ifile = '/tigress/wenchang/analysis/LMHU/data/Lizzie/sedimentHU_v20220711ens_noCaySal_wy_median_lp40_count.nc'
#da = xr.open_dataset(ifile)['sedimentHUcount'].sel(member='sediment_estimate')#.sel(year=slice(850, year_end))
da = xr.open_dataset(ifile)['sedimentHUcount'].sel(member='sediment_estimate')#.sel(year=slice(850, year_end))
da_lizzie = da
#upper and lower bounds
da = xr.open_dataset(ifile)['sedimentHUcount'].sel(member='lower_estimate')#.sel(year=slice(850, year_end))
da_lower = da
da = xr.open_dataset(ifile)['sedimentHUcount'].sel(member='upper_estimate')#.sel(year=slice(850, year_end))
da_upper = da
#the raw sediment HU
ifile = '/tigress/wenchang/analysis/LMHU/data/Lizzie/sedimentHU_v20220711ens_noCaySal_wy_max_lp40_count.nc'
da = xr.open_dataset(ifile)['sedimentHUcount'].sel(member='sediment_estimate')#.sel(year=slice(850, year_end))
da_lizzie_max = da

#sst HU
ifile = '/tigress/wenchang/analysis/LMHU/data/LMR/LMR2019_ssttc_rMDRtransformed18702000_HU_lp40_CI95.nc'
da = xr.open_dataset(ifile)['HU'].sel(quantile=0.5)#.sel(year=slice(850, year_end))
da_ssthu = da

#relative MDR SST
ifile = '/projects/w/wenchang/analysis/LastMillenniumHurricanes/LMHU/data/LMR/LMR2019_sstmdrtrop_rMDRtransformed18702000.nc'
ds = xr.open_dataset(ifile)
da = ds['MDR'].mean('MCrun') - ds['TROP'].mean('MCrun')
da_rMDR = da

#nino3.4
ifile = '/projects/w/wenchang/modelInput/LMR/LMR2019_nino34.nc'
da = xr.open_dataarray(ifile).mean('MCrun')
da_nino34 = da

def wyplot(ax=None, **kws):
    iframe = get_kws_from_argv('iframe', None)
    if iframe is not None: iframe = int(iframe)
    #SI = True if 'SI' in sys.argv else False #for supported info
    if ax is None: fig,ax = plt.subplots()
    year_start,year_end = 1350,1550

    #sediment records
    label = '(left) sediment HU'
    da = da_lizzie
    if iframe is None or iframe >= 1:
        ax.fill_between(da.year, da_lower.pipe(norm_bound), da_upper.pipe(norm_bound), color='gray', alpha=0.2)
        da.pipe(norm).pipe(lowpass).pipe(units).sel(year=slice(year_start, year_end)).plot(ax=ax, label=label, color='k')
        #da_lizzie_max.pipe(norm).pipe(lowpass).pipe(units).sel(year=slice(year_start, year_end)).plot(ax=ax, color='k', ls='--')
    print(label, da.pipe(norm).pipe(pct_diff).item())

    """
    #test: the raw HU, taking the max in the Monte Carlo bootstrap, overestimate the decline trend (-31% instead of -22%)
    label = 'sediment HU raw'
    da = da_lizzie_max 
    da.pipe(norm).pipe(lowpass).pipe(units).sel(year=slice(year_start, year_end)).plot(ax=ax, label=label, color='k', ls='--')
    print(label, da.pipe(norm).pipe(pct_diff).item())
    """
    
    #sst HU
    label = '(left) SST HU'
    da = da_ssthu
    if iframe is None or iframe >= 1:
        da.pipe(norm).pipe(lowpass).pipe(units).sel(year=slice(year_start, year_end)).plot(ax=ax, label=label, color='gray')
    print(label, da.pipe(norm).pipe(pct_diff).item())

    ax.axhline(0, color='gray', ls='--')
    ax.legend(loc='upper left')
    #title = f'{basin} HU, {n_window}-year-lowpass'
    title = f'{basin} HU and SST anomalies, {n_window}-year-lowpass'
    #if minWindmax == 33: title = title.replace('TC', 'HU')
    ax.set_title(title)
    ax.set_xlim(year_start, year_end)
    ax.axvspan(1350,1450, color='C3', alpha=0.1)
    ax.axvspan(1450,1550, color='C0', alpha=0.1)
    ax.set_ylim(-40, 40)

    ax2 = ax.twinx()
    #rMDR
    label = '(right) relative MDR SST'
    da = da_rMDR
    if iframe is None or iframe >= 1:
        da.pipe(anom).pipe(lowpass).sel(year=slice(year_start, year_end)).plot(ax=ax2, label=label)

    #nino34
    label = '(right) Nino3.4'
    da = da_nino34
    if iframe is None or iframe >= 1:
        da.pipe(anom).pipe(lowpass).sel(year=slice(year_start, year_end)).plot(ax=ax2, label=label, color='C2')

    ax2.set_ylim(-0.4, 0.4)
    ax2.legend(loc='upper right')
    ax2.grid(False)
    ax2.set_ylabel('SST anomaly [K]')

if __name__ == '__main__':
    from wyconfig import * #my plot settings
    wyplot()

    #savefig
    if 'savefig' in sys.argv or 's' in sys.argv:
        figname = __file__.replace('.py', f'.png')
        if 'overwritefig' in sys.argv or 'o' in sys.argv:
            wysavefig(figname, overwritefig=True)
        else:
            wysavefig(figname)
 
    try:
        tt.check(f'**Done**')
    except:
        pass
    print()
    if 'notshowfig' in sys.argv or 'n' in sys.argv:
        pass
    else:
        if 'plt' in globals(): plt.show()