#!/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,2024)
#years = slice(1991,2020)
years = slice(1991,2026)
#hurdat2
"""
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
"""
ifile = '/projects/w/wenchang/analysis/active_work/hurdat2/nHU.hurdat2.1851-2024.nc'
da = xr.open_dataarray(ifile).sel(year=years)
#2025: https://en.wikipedia.org/wiki/2025_Atlantic_hurricane_season
da2025 = xr.DataArray([5,], dims=('year',), coords=([2025,],))
da = xr.concat([da, da2025], 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 v04r01_ts33 or HU
ds = xr.open_dataset('/projects/w/wenchang/data/ibtracs/v04r01/analysis/IBTrACS.ALL.v04r01.2024-10-24.counts.TS33.1980-2023.yearly.nc')
ds_v04r01_ts33 = 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.TS33.5ens.1871-2021.yearly.nc',
    '/projects/w/wenchang/analysis/TC/HIRAM/amipHadISSTlongChancorr_tigercpu_intelmpi_18_540PE_extend2019-2021/netcdf/tc_counts.TS33.5ens.2022-2023.yearly.nc',
    '/projects/w/wenchang/analysis/TC/HIRAM/amipHadISSTlongChancorr_tigercpu_intelmpi_18_540PE_extend2019-2021/netcdf/tc_counts.TS33.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.TS33.10ens.1871-2021.yearly.nc',
    '/projects/w/wenchang/analysis/TC/AM2.5C360/amipHadISSTlong_chancorr_tigercpu_intelmpi_18_1080PE/netcdf/tc_counts.TS33.10ens.2022-2022.yearly.nc',
    '/projects/w/wenchang/analysis/TC/AM2.5C360/amipHadISSTlong_chancorr_tigercpu_intelmpi_18_1080PE/netcdf/tc_counts.TS33.10ens.2023-2023.yearly.nc',
    '/projects/w/wenchang/analysis/TC/AM2.5C360/amipHadISSTlong_chancorr_tigercpu_intelmpi_18_1080PE/netcdf/tc_counts.TS33.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.TS33.5ens.1871-2020.yearly.nc
/projects/w/wenchang/analysis/TC/AM2.5C360/amipHadISSTlong_chancorr_rmWarming_tigercpu_intelmpi_18_1080PE/netcdf/tc_counts.TS33.5ens.2021-2023.yearly.nc
/projects/w/wenchang/analysis/TC/AM2.5C360/amipHadISSTlong_chancorr_rmWarming_tigercpu_intelmpi_18_1080PE/netcdf/tc_counts.TS33.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

#nmme1997
ifile = '/projects/w/wenchang/analysis/TC/AM2.5C360/nmmeHindcast199701_tiger3_intel24ifort_openmpi_2430PE/netcdf/tc_counts.TS33.10ens.1997-1997.yearly.nc'
ds97 = xr.open_dataset(ifile)
#nmme1998
ifile = '/projects/w/wenchang/analysis/TC/AM2.5C360/nmmeHindcast199801_tiger3_intel24ifort_openmpi_2430PE/netcdf/tc_counts.TS33.10ens.1998-1998.yearly.nc'
ds98 = xr.open_dataset(ifile)
#nmme1991-2020
ifile = '/projects/w/wenchang/analysis/TC/AM2.5C360/nmmeHindcastJan01L12mon_tiger3_intel24ifort_openmpi_2430PE/netcdf/tc_counts.TS33.10ens.1991-2020.yearly.nc'
ds_nmme = xr.open_dataset(ifile)
#nmme201005
ifile = '/projects/w/wenchang/analysis/TC/AM2.5C360/nmmeHindcast201005_tiger3_intel24ifort_openmpi_2430PE/netcdf/tc_counts.TS33.5ens.2010-2010.yearly.nc'
ds2010may = xr.open_dataset(ifile)
#nmmeForecast202602
ifile = '/projects/w/wenchang/analysis/TC/AM2.5C360/nmmeForecast202602_tiger3_intel24ifort_openmpi_2430PE/netcdf/tc_counts.TS33.10ens.2026-2026.yearly.nc'
ds2026feb = xr.open_dataset(ifile)
#nmmeMay1991-2020
ifile = '/projects/w/wenchang/analysis/TC/AM2.5C360/nmmeHindcastMay01L12mon_tiger3_intel24ifort_openmpi_2430PE/netcdf/tc_counts.TS33.10ens.1991-2020.yearly.nc'
ds_nmmeMay = xr.open_dataset(ifile)
#nmmeForecast202605
ifile = '/projects/w/wenchang/analysis/TC/AM2.5C360/nmmeForecast202605_tiger3_intel24ifort_openmpi_2430PE/netcdf/tc_counts.TS33.10ens.2026-2026.yearly.nc'
ds2026may = xr.open_dataset(ifile)

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
    import xlinregress
    hiram_only = False #default
    mulAdjust = False #True #default
    ibtracsOn = False #True
    basin = 'NA'
    da = da_hurdat2
    da.plot(color='k', lw=2, label='HURDAT2', marker='o', fillstyle='none')
    #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 ibtracsOn:
        ds_v04r01_ts33[basin].plot(color='gray', lw=2, label='IBTrACS_v04r01_TS33')
    
    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 AMIP', 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='--')
        #nmme1997
        da = ds97[basin]
        for ii in range(da.en.size):
            if ii==0:
                plt.plot(da.year, da.isel(en=ii), marker='o', fillstyle='none', color='C1', label='hindcast1997', ls='')
            else:
                plt.plot(da.year, da.isel(en=ii), marker='o', fillstyle='none', color='C1')
        plt.plot(da.year, da.mean('en'), marker='x', color='C1')
        #nmme1998
        da = ds98[basin]
        for ii in range(da.en.size):
            if ii==0:
                plt.plot(da.year, da.isel(en=ii), marker='o', fillstyle='none', color='C2', label='hindcast1998', ls='')
            else:
                plt.plot(da.year, da.isel(en=ii), marker='o', fillstyle='none', color='C2')
        plt.plot(da.year, da.mean('en'), marker='x', color='C2')
        """
        #nmme1991-2020
        da = ds_nmme[basin]
        da.mean('en').plot(color='C3', label='AM2.5C360 NMME Jan', ls='--')
        for ii in range(da.en.size):
            da.isel(en=ii).plot(color='C3', ls='--', lw=1, alpha=0.3)
        #corr
        ys = slice(1991, 2020)
        xx = da_hurdat2.sel(year=ys)
        yy = ds_nmme[basin].mean('en').sel(year=ys)
        rg = yy.linregress.on(xx, dim='year')
        print('nmme Jan and hurdat2', rg.r.values)
        xx = ds_am2p5c360[basin].sel(year=ys).mean('en')
        rg = yy.linregress.on(xx, dim='year')
        print('nmme Jan and AMIP', rg.r.values)
        """
        #nmme201005
        da = ds2010may[basin]
        for ii in range(da.en.size):
            if ii==0:
                plt.plot(da.year, da.isel(en=ii), marker='o', fillstyle='none', color='C4', label='hindcast2010May', ls='')
            else:
                plt.plot(da.year, da.isel(en=ii), marker='o', fillstyle='none', color='C4')
        plt.plot(da.year, da.mean('en'), marker='x', color='C4')
        """
        #nmmeMay1991-2020
        da = ds_nmmeMay[basin]
        da.mean('en').plot(color='C2', label='AM2.5C360 NMME May', ls='--')
        for ii in range(da.en.size):
            da.isel(en=ii).plot(color='C2', ls='--', lw=1, alpha=0.3)
        #corr
        ys = slice(1991, 2020)
        xx = da_hurdat2.sel(year=ys)
        yy = ds_nmmeMay[basin].mean('en').sel(year=ys)
        rg = yy.linregress.on(xx, dim='year')
        print('nmme May and hurdat2', rg.r.values)
        xx = ds_am2p5c360[basin].sel(year=ys).mean('en')
        rg = yy.linregress.on(xx, dim='year')
        print('nmme May and AMIP', rg.r.values)
            
        #nmmeForecast202602
        da = ds2026feb[basin]
        offset = -0.2
        color = 'C0'
        alpha = 0.5
        for ii in range(da.en.size):
            if ii==0:
                plt.plot(da.year+offset, da.isel(en=ii), marker='o', fillstyle='none', color=color, label='forecast2026Feb', ls='', alpha=alpha)
            else:
                plt.plot(da.year+offset+0.05*ii, da.isel(en=ii), marker='o', fillstyle='none', color=color, alpha=alpha)
        plt.plot(da.year+offset, da.mean('en'), marker='x', color=color)

        #nmmeForecast202605
        da = ds2026may[basin]
        offset = 0.2
        color = 'C1'
        alpha = 0.5
        for ii in range(da.en.size):
            if ii==0:
                plt.plot(da.year+offset, da.isel(en=ii), marker='o', fillstyle='none', color=color, label='forecast2026May', ls='', alpha=alpha)
            else:
                plt.plot(da.year+offset+0.05*ii, da.isel(en=ii), marker='o', fillstyle='none', color=color, alpha=alpha)
        plt.plot(da.year+offset, da.mean('en'), marker='x', color=color)

    """
    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(ncol=2, loc='upper left')
    ax.set_title(f'{basin} HU counts from models and obs.')
    ax.set_ylabel('#')
    from matplotlib.ticker import MaxNLocator
    ax.xaxis.set_major_locator(MaxNLocator(integer=True))
    ax.set_ylim(None, 20)
    #ax.axvspan(2010-0.5, 2010+0.5, color='gray', alpha=0.2)

    #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()