#!/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 misc.seasons import sel_season
import xfilter
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
iframe = get_kws_from_argv('iframe', None)
if iframe is not None: iframe = int(iframe)
year_start,year_end = 1350,1550
basin = """
global
WP
NA
""".split()[-1]
#timeslice_ref = slice('1900', '1999')
#yearslice_ref = slice(1900, 1999)
timeslice_ref = slice('1351', '1550')
yearslice_ref = slice(1351, 1550)
#timeslice_ref = slice('900', '1999')
#yearslice_ref = slice(900, 1999)
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
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
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
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
##from ensemble
#ifile = '/tigress/wenchang/analysis/LastMillenniumHurricanes/LMHU/data/Lizzie/sedimentHU_v20220711ens_noCaySal.nc'
#da = xr.open_dataarray(ifile).isel(jk=0).quantile([0.1, 0.5, 0.9])
#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
#HiRAM HU
ifile = '/projects/w/wenchang/analysis/TC/HIRAM/amipLMR2019SST0850ic_tigercpu_intelmpi_18_540PE/netcdf/tc_counts.TS30.ens01.0850-1999.nc'
ds_hiram_hu = xr.open_dataset(ifile)
"""
#HiRAM TC 
ifile = '/projects/w/wenchang/analysis/TC/HIRAM/amipLMR2019SST0850ic_tigercpu_intelmpi_18_540PE/netcdf/tc_counts.TS17.ens01.0850-1999.nc'
ds_hiram_tc = xr.open_dataset(ifile)
#HiRAM seeds
ifile = '/projects/w/wenchang/analysis/TC/HIRAM/amipLMR2019SST0850ic_tigercpu_intelmpi_18_540PE/netcdf/tc_counts.seeds.ens01.0850-1999.minohour12.nc'
ds_hiram_sd = xr.open_dataset(ifile)
ifile_seeds = ifile #save data for later use

#HiRAM transition probability
ifile = 'prob_HIRAM_amipLMR2019SST0850ic_tigercpu_intelmpi_18_540PE_ens01_0850-1999_NAmean.nc'
da = xr.open_dataarray(ifile)#.groupby('time.year').mean('time')
da_prob = da ** 2 #adjust for HU from TC
ifile_prob = ifile #save data for later use

#HiRAM nseeds x prob
time_prob = da_prob.time
da_nseedsXprob = ds_hiram_sd[basin].assign_coords(time=time_prob) * da_prob
 
"""
#spi
ifile = 'spi_HIRAM_amipLMR2019SST0850ic_tigercpu_intelmpi_18_540PE_ens01_0850-1999_NAmean.nc'
da = xr.open_dataarray(ifile)#.groupby('time.year').mean('time')
da_spi = da #adjust for HU from TC

#pvi
ifile = 'pvi_HIRAM_amipLMR2019SST0850ic_tigercpu_intelmpi_18_540PE_ens01_0850-1999_NAmean.nc'
da = xr.open_dataarray(ifile)#.groupby('time.year').mean('time')
da_pvi = da ** 2 #adjust for HU from TC

#spiXpvi
da = da_spi * da_pvi
da_spipvi = da

#HiRAM transition probability from PI
ifile = 'pviPI_HIRAM_amipLMR2019SST0850ic_tigercpu_intelmpi_18_540PE_ens01_0850-1999_NAmean.nc'
da = xr.open_dataarray(ifile)#.groupby('time.year').mean('time')
da_pviPI = da**2
#HiRAM transition probability from Vshear
ifile = 'pviVshear_HIRAM_amipLMR2019SST0850ic_tigercpu_intelmpi_18_540PE_ens01_0850-1999_NAmean.nc'
da = xr.open_dataarray(ifile)#.groupby('time.year').mean('time')
da_pviVshear = da**2
#HiRAM transition probability from chi
ifile = 'pvichi_HIRAM_amipLMR2019SST0850ic_tigercpu_intelmpi_18_540PE_ens01_0850-1999_NAmean.nc'
da = xr.open_dataarray(ifile)#.groupby('time.year').mean('time')
da_pvichi = da**2

def wyplot(ax=None, **kws):
    if ax is None: fig,ax = plt.subplots()
    #if minWindmax in (17, 33):
    #year_start,year_end = 1250,1550
    #year_start,year_end = 850,1999

    #sediment records
    label = '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())

    #sst HU
    label = '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())

    #HiRAM HU
    label = 'HiRAM HU'
    da = ds_hiram_hu[basin].groupby('time.year').sum('time')
    da_ = da.pipe(norm).pipe(lowpass).pipe(units)
    if iframe is None or iframe >= 2:
        da_.plot(label=label, ax=ax, color='C3')
    print(label, da.pipe(norm).pipe(pct_diff).item())

    #HiRAM spi * pvi
    label = r'SPI$\times$p(VI)'
    da = da_spipvi.groupby('time.year').sum('time')
    da_ = da.pipe(norm).pipe(lowpass).pipe(units)
    #if iframe is None or iframe >= 5:
    #    da_.plot(label=label, ax=ax, color='C2')
    da_.plot(label=label, ax=ax, color='C2')
    #da_ = _da_seeds + _da_prob #get from saved data before
    #da_.isel(year=slice(n_window//2, -n_window//2)).plot(ax=ax, color='C2', ls='--')
    print(label, da.pipe(norm).pipe(pct_diff).item())

    #spi
    label = 'SPI'
    da = da_spi.groupby('time.month') * da_pvi.groupby('time.month').mean('time')
    da = da.groupby('time.year').sum('time')
    da_ = da.pipe(norm).pipe(lowpass).pipe(units)
    _da_seeds = da_ #save for later use
    if iframe is None or iframe >= 3:
        da_.plot(label=label, ax=ax, color='C0', ls='--')
    print(label, da.pipe(norm).pipe(pct_diff).item())

    #pvi
    label = 'p(VI)'
    da = da_spi.groupby('time.month').mean('time') * da_pvi.groupby('time.month')
    da = da.groupby('time.year').mean('time')
    da_ = da.pipe(norm).pipe(lowpass).pipe(units)
    _da_prob = da_ #save for later use
    if iframe is None or iframe >= 4:
        da_.plot(label=label, ax=ax, color='C1', ls='--')
    print(label, da.pipe(norm).pipe(pct_diff).item())

    #HiRAM pvi from PI
    label = 'p(VI) from PI'
    da = da_spi.groupby('time.month').mean('time') * da_pviPI.groupby('time.month')
    da = da.groupby('time.year').mean('time')
    da_ = da.pipe(norm).pipe(lowpass).pipe(units)
    #da_.plot(label=label, ax=ax, color='C1', ls=':')
    print(label, da.pipe(norm).pipe(pct_diff).item())

    #HiRAM pvi from Vshear
    label = 'p(VI) from wind shear'
    da = da_spi.groupby('time.month').mean('time') * da_pviVshear.groupby('time.month')
    da = da.groupby('time.year').mean('time')
    da_ = da.pipe(norm).pipe(lowpass).pipe(units)
    if iframe is None or iframe >= 5:
        da_.plot(label=label, ax=ax, color='C1', ls=':')
    print(label, da.pipe(norm).pipe(pct_diff).item())

    #HiRAM pvi from chi
    label = 'p(VI) from $\chi$'
    da = da_spi.groupby('time.month').mean('time') * da_pvichi.groupby('time.month')
    da = da.groupby('time.year').mean('time')
    da_ = da.pipe(norm).pipe(lowpass).pipe(units)
    #da_.plot(label=label, ax=ax, color='C1', ls=':')
    print(label, da.pipe(norm).pipe(pct_diff).item())
    """
    #HiRAM TC
    label = 'HiRAM TC'
    ds = ds_hiram_tc.groupby('time.year').sum('time')
    da_ = ds[basin].pipe(norm).pipe(lowpass).pipe(units)
    da_.plot(label=label, ax=ax, color='C1')

    """

    ax.axhline(0, color='gray', ls='--')
    ax.legend(loc='upper right', ncol=1)
    title = f'{basin} HU, {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(-30, 40)

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'__{n_window}lowpass_{year_start:04d}-{year_end:04d}.png')
        if iframe is not None: 
            figname = figname.replace(os.path.basename(figname), 'frames/'+os.path.basename(figname))
            figname = figname.replace('.png', f'_frame{iframe:02d}.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()