#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Wed Nov  6 11:08:24 EST 2024
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
#
if __name__ == '__main__':
    try:
        tt.check('end import')
    except:
        pass
#
#start from here
#years = slice(1980,2023)
years = slice(1980,2024)
#hurdat
df = pd.read_csv('https://tigress-web.princeton.edu/~gvecchi/hurdat2_long_adjusted.txt', header=None, sep='\s+')
da = xr.DataArray(df.iloc[:,1].values, dims=('year',), coords=(df.iloc[:,0].astype('int'),)).sel(year=years)
da2024 = xr.DataArray([14,], dims=('year',), coords=([2024,],))
da = xr.concat([da, da2024], dim='year')
da_hurdat2 = da
#ibtracs v04r01_ts17_2days
ds = xr.open_dataset('/projects/w/wenchang/data/ibtracs/v04r01/analysis/IBTrACS.ALL.v04r01.2024-10-24.counts.TS17.2days17.1980-2023.yearly.nc')
ds_v04r01_ts17_2days = ds
#ibtracs v04r01_ts17
ds = xr.open_dataset('/projects/w/wenchang/data/ibtracs/v04r01/analysis/IBTrACS.ALL.v04r01.2024-10-24.counts.TS17.1980-2023.yearly.nc')
ds_v04r01_ts17 = ds
"""
#ibtracs v04r00_ts17_2days
ds = xr.open_dataset('/projects/w/wenchang/data/ibtracs/v04r00/analysis/v2/IBTrACS.ALL.v04r00.2022-01-25.counts.TS17.2days17.1980-2021.yearly.nc')
ds_v04r00_ts17_2days = ds
#ibtracs v04r00_ts17
ds = xr.open_dataset('/projects/w/wenchang/data/ibtracs/v04r00/analysis/v2/IBTrACS.ALL.v04r00.2022-01-25.counts.TS17.1980-2021.yearly.nc')
ds_v04r00_ts17 = ds
"""

#hiram
ifiles = [
    '/projects/w/wenchang/analysis/TC/HIRAM/amipHadISSTlongChancorr_tigercpu_intelmpi_18_540PE_extend2019-2021/netcdf/tc_counts.TS17.5ens.1871-2021.yearly.nc',
    '/projects/w/wenchang/analysis/TC/HIRAM/amipHadISSTlongChancorr_tigercpu_intelmpi_18_540PE_extend2019-2021/netcdf/tc_counts.TS17.5ens.2022-2023.yearly.nc',
    '/projects/w/wenchang/analysis/TC/HIRAM/amipHadISSTlongChancorr_tigercpu_intelmpi_18_540PE_extend2019-2021/netcdf/tc_counts.TS17.5ens.2024-2024.yearly.nc'
    ]
ds = xr.open_mfdataset(ifiles).load().sel(year=years)
#dss = []
#for ifile in ifiles:
#    dss.append(xr.open_dataset(ifile))
#ds = xr.concat(dss, dim='year')
ds_hiram = ds

#am2.5c360
ifiles = [
    '/projects/w/wenchang/analysis/TC/AM2.5C360/amipHadISSTlong_chancorr_tigercpu_intelmpi_18_1080PE/netcdf/tc_counts.TS17.10ens.1871-2021.yearly.nc',
    '/projects/w/wenchang/analysis/TC/AM2.5C360/amipHadISSTlong_chancorr_tigercpu_intelmpi_18_1080PE/netcdf/tc_counts.TS17.10ens.2022-2022.yearly.nc',
    '/projects/w/wenchang/analysis/TC/AM2.5C360/amipHadISSTlong_chancorr_tigercpu_intelmpi_18_1080PE/netcdf/tc_counts.TS17.10ens.2023-2023.yearly.nc',
    '/projects/w/wenchang/analysis/TC/AM2.5C360/amipHadISSTlong_chancorr_tigercpu_intelmpi_18_1080PE/netcdf/tc_counts.TS17.10ens.2024-2024.yearly.nc'
    ]
ds = xr.open_mfdataset(ifiles).load().sel(year=years)
#dss = []
#for ifile in ifiles:
#    dss.append(xr.open_dataset(ifile))
#ds = xr.concat(dss, dim='year')
ds_am2p5c360 = ds

#am2.5c360 rmWarming
ifiles = """
/projects/w/wenchang/analysis/TC/AM2.5C360/amipHadISSTlong_chancorr_rmWarming_tigercpu_intelmpi_18_1080PE/netcdf/tc_counts.TS17.5ens.1871-2020.yearly.nc
/projects/w/wenchang/analysis/TC/AM2.5C360/amipHadISSTlong_chancorr_rmWarming_tigercpu_intelmpi_18_1080PE/netcdf/tc_counts.TS17.5ens.2021-2023.yearly.nc
/projects/w/wenchang/analysis/TC/AM2.5C360/amipHadISSTlong_chancorr_rmWarming_tigercpu_intelmpi_18_1080PE/netcdf/tc_counts.TS17.5ens.2024-2024.yearly.nc
""".split()
ds = xr.open_mfdataset(ifiles).load().sel(year=years)
ds_am2p5c360_nw = ds #no warming

#ERA5 spiXpvi
ifile = '/tigress/wenchang/analysis/active_work/histTC/ERA5/spiXpvi_ERA5_1979-2024_NAmean.nc'
da = xr.open_dataarray(ifile).groupby('time.year').mean('time')
da_era5 = da

def mul_adjust(da, da_norm=None):
    if da_norm is None: da_norm = da
    #years_ref = slice(1982,2023)
    years_ref = slice(1981,2010)
    if 'en' in da.dims:
        return da * da_hurdat2.sel(year=years_ref).mean('year')/da_norm.sel(year=years_ref).mean('year').mean('en')
    else:
        return da * da_hurdat2.sel(year=years_ref).mean('year')/da_norm.sel(year=years_ref).mean('year')
 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    hiram_only = False #default
    mulAdjust = True #default
    ibtracsOn = False
    basin = 'NA'
    da = da_hurdat2
    da.plot(color='k', lw=2, label='HURDAT2 corrected', marker='o', fillstyle='none')
    if ibtracsOn:
        ds_v04r01_ts17_2days[basin].plot(color='gray', lw=2, label='IBTrACS_v04r01_TS17_2days')
    #ds_v04r01_ts17[basin].plot(color='gray', lw=2, label='IBTrACS_v04r01_TS17', ls='--')
    #ds_v04r00_ts17_2days[basin].plot(color='k', lw=2, label='IBTrACS_v04r00_TS17_2days', ls='--')
    #ds_v04r00_ts17[basin].plot(color='gray', lw=2, label='IBTrACS_v04r00_TS17', ls='--')
    
    if not hiram_only:
        da = ds_am2p5c360[basin]
        if mulAdjust: da = mul_adjust(da)
        #da.plot(hue='en', color='C0',  lw=1, alpha=0.3, add_legend=False, ls='-')
        plt.fill_between(da.year, da.min('en'), da.max('en'), color='C0', alpha=0.2)
        da.mean('en').plot(color='C0', label='AM2.5C360', ls='-')
        #no warming
        da = ds_am2p5c360_nw[basin]
        if mulAdjust: da = mul_adjust(da, ds_am2p5c360[basin])
        #da.plot(hue='en', color='C0',  lw=1, alpha=0.3, add_legend=False, ls='-')
        da.mean('en').plot(color='C0', label='AM2.5C360 rmWarming', ls='--')

    da = ds_hiram[basin]
    if mulAdjust: da = mul_adjust(da)
    #da.plot(hue='en', color='C1',  lw=1, alpha=0.3, add_legend=False, ls='-')
    plt.fill_between(da.year, da.min('en'), da.max('en'), color='C1', alpha=0.2)
    da.mean('en').plot(color='C1', label='HiRAM', ls='-')
    
    if not hiram_only:
        #ERA5 spi x pvi
        da = da_era5
        da = mul_adjust(da)
        da.plot(color='C2', label=r'ERA5 SPI$\times$p(VI)', ls='-')
    ax = plt.gca()
    ax.legend()
    ax.set_title(f'{basin} TC counts from models and obs.')
    ax.set_ylabel('#')

    #savefig
    if 'savefig' in sys.argv or 's' in sys.argv:
        figname = __file__.replace('.py', f'.png')
        if hiram_only: figname = figname.replace('.png', '__hiramOnly.png')
        if mulAdjust: figname = figname.replace('.png', '__mulAdjust.png')
        if ibtracsOn: figname = figname.replace('.png', '__withIBTrACS.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()