#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Fri Sep 15 09:51:11 EDT 2023
if __name__ == '__main__':
    import sys,os
    from misc.timer import Timer
    tt = Timer(f'[{os.getcwd()}] start ' + ' '.join(sys.argv))
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
##old way
#from tc_monitor_AM2p5_amipLMR2019SST0850ic_tigercpu_intelmpi_18_540PE import ds_counts_yearly as ds_am2p5
#from tc_monitor_HIRAM_amipLMR2019SST0850ic_tigercpu_intelmpi_18_540PE import ds_counts_yearly as ds_hiram
from modelout.getTC import getTC
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
basin = """
global
WP
NA
""".split()[-1]
minWindmax = 33 if '33' in sys.argv else 17 #default is 17m/s
if 'HU' in sys.argv:
    tc_type = 'HU'
    windmax2min_am2p5 = 27
    windmax2min_hiram = 30
else:
    tc_type = 'TC'
    windmax2min_am2p5 = 17
    windmax2min_hiram = 17
ens = 1
model = 'AM2.5'
expname = 'amipLMR2019SST0850ic_tigercpu_intelmpi_18_540PE'
ds_am2p5 = getTC(model=model, expname=expname, ens=ens, minWindmax=windmax2min_am2p5, storm_type='allstorms', dHour=-24, basin='NA')
model = 'HIRAM'
expname = 'amipLMR2019SST0850ic_tigercpu_intelmpi_18_540PE'
ds_hiram = getTC(model=model, expname=expname, ens=ens, minWindmax=windmax2min_hiram, basin='NA')
 
 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    import xfilter
    norm = lambda da: da/da.sel(year=slice(850,950)).mean('year')*100 - 100
    units = lambda da: da.assign_attrs(units='% of 850-950 mean', long_name='counts')
    lowpass  = lambda da: da.rolling(year=9, min_periods=1, center=True).mean()
    fig,ax = plt.subplots()
    alpha = 0.5
    ds = ds_am2p5
    label = 'AM2.5'
    ds[basin].pipe(norm).pipe(units).plot(ax=ax, color='C0', lw=1, alpha=alpha)
    ds[basin].pipe(norm).pipe(lowpass).pipe(units).plot(label=label, ax=ax, color='C0')
    ds = ds_hiram
    label = 'HiRAM'
    ds[basin].pipe(norm).pipe(units).plot(ax=ax, color='C1', lw=1, alpha=alpha)
    ds[basin].pipe(norm).pipe(lowpass).pipe(units).plot(label=label, ax=ax, color='C1')

    ax.axhline(0, color='gray', ls='--')
    ax.legend(loc='upper left')
    title = f'{basin} {tc_type} #'
    ax.set_title(title)

    #savefig
    if 'savefig' in sys.argv or 's' in sys.argv:
        figname = __file__.replace('.py', f'__{tc_type}.png')
        if 'overwritefig' in sys.argv or 'o' in sys.argv:
            wysavefig(figname, overwritefig=True)
        else:
            wysavefig(figname)

    #40-year lowpass
    n_window = 40 #40 years
    lowpass = lambda da: da.filter.lowpass(1/n_window, dim='year', padtype='even')
    fig,ax = plt.subplots()
    year_end = max(ds_am2p5.year[-1].item(), ds_hiram.year[-1].item()) #year max for both AM2.5 and HiRAM simulations
    year_end = min(year_end+100, 2000) #100 more years but not exceed 2000
    #if minWindmax in (17, 33):
    if tc_type in ('TC', 'HU'):
        #sediment records
        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.pipe(norm).plot(ax=ax, label='sedimentHU', color='k')
        #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.pipe(norm).plot(ax=ax, label='sstHU', color='gray')

    label = 'AM2.5'
    ds = ds_am2p5
    da_ = ds[basin].pipe(norm).pipe(lowpass).pipe(units)
    da_am2p5 = da_#cached
    da_.isel(year=slice(n_window//2, -n_window//2)).plot(label=label, ax=ax, color='C0')
    da_.isel(year=slice(0, n_window//2)).plot(ax=ax, color='C0', ls='--')
    da_.isel(year=slice(-n_window//2, None)).plot(ax=ax, color='C0', ls='--')
    label = 'HiRAM'
    ds = ds_hiram
    da_ = ds[basin].pipe(norm).pipe(lowpass).pipe(units)
    da_hiram = da_ #cached
    da_.isel(year=slice(n_window//2, -n_window//2)).plot(label=label, ax=ax, color='C1')
    da_.isel(year=slice(0, n_window//2)).plot(ax=ax, color='C1', ls='--')
    da_.isel(year=slice(-n_window//2, None)).plot(ax=ax, color='C1', ls='--')
    label = '(AM2.5+HiRAM)/2'
    da_ = (da_am2p5 + da_hiram)/2
    da_.pipe(units).plot(label=label, ax=ax, color='C2', ls='--')

    ax.axhline(0, color='gray', ls='--')
    ax.legend(loc='upper left', ncol=2)
    title = f'{basin} {tc_type} #, {n_window}-lowpass'
    #if minWindmax == 33: title = title.replace('TC', 'HU')
    ax.set_title(title)
    ax.set_xlim(None, year_end)

    #savefig
    if 'savefig' in sys.argv or 's' in sys.argv:
        figname = __file__.replace('.py', f'__{n_window}lowpass_{tc_type}.png')
        if 'overwritefig' in sys.argv or 'o' in sys.argv:
            wysavefig(figname, overwritefig=True)
        else:
            wysavefig(figname)

    tt.check(f'**Done**')
    print()
    if 'notshowfig' in sys.argv or 'n' in sys.argv:
        pass
    else:
        if 'plt' in globals(): plt.show()