#!/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
basin = 'NA'

#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_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
#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).rename(en='ens')
#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).rename(en='ens')

##HiRAM SPI
#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).rename(en='ens')
#ifile_seeds = ifile #save data for later use
#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

#HiRAM transition probability
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
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
#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

#percent decline from 1351-1450 to 1451-1550
if True:
    dc = {}
    #timeslice_ref = slice('1900', '1999')
    #yearslice_ref = slice(1900, 1999)
    timeslice_ref = slice('1351', '1550')
    yearslice_ref = slice(1351, 1550)
    norm = lambda da: da/da.sel(year=yearslice_ref).mean('year')*100 - 100
    pct_diff = lambda da: da.sel(year=slice(1451, 1550)).mean('year') - da.sel(year=slice(1351,1450)).mean('year')
    label = 'sediment'
    da = da_lizzie
    da = da.pipe(norm).pipe(pct_diff)
    dc[label] = da

    label = 'SST'
    da = da_ssthu
    da = da.pipe(norm).pipe(pct_diff)
    dc[label] = da

    label = 'HiRAM'
    da = ds_hiram_hu[basin].groupby('time.year').sum('time')
    da = da.pipe(norm).pipe(pct_diff)
    dc[label] = da

    """
    label = 'seeds and prob'
    da = da_nseedsXprob.groupby('time.year').sum('time')
    da = da.pipe(norm).pipe(pct_diff)
    dc[label] = da
    """
    label = r'SPI$\times$p(VI)'
    da = da_spipvi.groupby('time.year').sum('time')
    da = da.pipe(norm).pipe(pct_diff)
    dc[label] = da

    label = 'SPI' #'seeds'
    #da = ds_hiram_sd[basin].groupby('time.month') * da_prob.sel(time=timeslice_ref).groupby('time.month').mean('time')
    da = da_spi.groupby('time.month') * da_pvi.sel(time=timeslice_ref).groupby('time.month').mean('time')
    da = da.groupby('time.year').mean('time')
    da = da.pipe(norm).pipe(pct_diff)
    dc[label] = da

    label = 'p(VI)' #'prob'
    #da = ds_hiram_sd[basin].sel(time=timeslice_ref).groupby('time.month').mean('time') * da_prob.groupby('time.month')
    da = da_spi.sel(time=timeslice_ref).groupby('time.month').mean('time') * da_pvi.groupby('time.month')
    da = da.groupby('time.year').mean('time')
    da = da.pipe(norm).pipe(pct_diff)
    dc[label] = da

    label = 'p(VI) from shear'
    da = da_spi.sel(time=timeslice_ref).groupby('time.month').mean('time') * da_pviVshear.groupby('time.month')
    da = da.groupby('time.year').mean('time')
    da = da.pipe(norm).pipe(pct_diff)
    dc[label] = da

    label = 'p(VI) from PI'
    da = da_spi.sel(time=timeslice_ref).groupby('time.month').mean('time') * da_pviPI.groupby('time.month')
    da = da.groupby('time.year').mean('time')
    da = da.pipe(norm).pipe(pct_diff)
    dc[label] = da

    label = 'p(VI) from $\chi$'
    da = da_spi.sel(time=timeslice_ref).groupby('time.month').mean('time') * da_pvichi.groupby('time.month')
    da = da.groupby('time.year').mean('time')
    da = da.pipe(norm).pipe(pct_diff)
    dc[label] = da

    ds = xr.Dataset(dc)

iframe = get_kws_from_argv('iframe', None)
if iframe is not None: iframe = int(iframe)

def wyplot(ax=None, **kws):
    if ax is None: fig, ax = plt.subplots()
    df = ds.mean('ens').to_pandas()
    colors = ['k', 'gray', 'C3', 'C2', 'C0'] +  ['C1',]*4
    df.plot.bar(rot=30, ax=ax, color=colors, **kws)
    ax.set_xticklabels(ax.get_xticklabels(), ha='right', rotation=30)
    #if iframe is not None:
    #    df.where(np.arange(df.size)<iframe).plot.bar(rot=30, ax=ax)
    ax.bar_label(ax.containers[0], fmt='%.2g')
        

    ax.set_ylabel('%')
    ax.set_title('Hurricane frequency change from 1351-1450 to 1451-1550')
    ax.axhline(0, color='gray', ls='--')
    ax.set_ylim(-25, 0)
    color, alpha = 'C1', 0.2
    if iframe == 2: ax.axvspan(0-0.5, 2+0.5, color=color, alpha=alpha)
    if iframe == 3: ax.axvspan(2-0.5, 4+0.5, color=color, alpha=alpha)
    if iframe == 4: ax.axvspan(4-0.5, 7+0.5, color=color, alpha=alpha)

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'__pctChange.png')
        if iframe is not None:
            basename = os.path.basename(figname)
            figname = figname.replace(basename, 'frames/'+basename)
            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()