#!/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
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
#model = 'HiRAM'
model = 'AM2.5C360'
ii = get_kws_from_argv('ens', None)
#ens = range(1,3+1)
if ii is None:
    ens = [16, 17, 18]
else:
    ii = int(ii)
    ens = [ii,]
basin = get_kws_from_argv('basin', 'NA')
#timeslice_ref = slice('1900', '1999')
#yearslice_ref = slice(1900, 1999)
timeslice_ref = slice('1991', '2020')
yearslice_ref = slice(1991, 2020)
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='NTC')
#pct_diff = lambda da: da.sel(year=slice(1451, 1550)).mean('year') - da.sel(year=slice(1351,1450)).mean('year')
#lowpass
n_window = 15 #
n_window = get_kws_from_argv('window', n_window)
n_window = int(n_window)
if 'rolling' in sys.argv:
    lowpass = lambda da: da.rolling(year=10, center=True).mean()
else:
    lowpass = lambda da: da.filter.lowpass(1/n_window, dim='year', padtype='even')

#NTC
#ifile = 'tc_counts.TS17.6ens.1871-2100.nc'
ifile = glob.glob('tc_counts.TS17.*.nc')[0]
ds = xr.open_dataset(ifile).rename(en='ens')
ntc = ds[basin].sel(ens=ens)
time_ntc = ntc.time.values
#spi
#ifile = 'spi_AM2.5C360_amipHadISSTrcp45_tigercpu_intelmpi_18_1080PE_ens16_1871-2100_NAmean.nc'
#ifile = glob.glob(f'spi_*_{basin}mean.nc')[0]
#da = xr.open_dataarray(ifile).assign_coords(time=time_ntc).sel(ens=ens)
ifiles = glob.glob(f'spi_*_{basin}mean.nc')
da = xr.open_mfdataset(ifiles)['spi'].assign_coords(time=time_ntc).sel(ens=ens)
spi = da
#transition probability
#ifile = 'pvi_AM2.5C360_amipHadISSTrcp45_tigercpu_intelmpi_18_1080PE_ens16_1871-2100_NAmean.nc'
#ifile = glob.glob(f'pvi_*_{basin}mean.nc')[0]
#da = xr.open_dataarray(ifile).assign_coords(time=time_ntc).sel(ens=ens)
ifiles = glob.glob(f'pvi_*_{basin}mean.nc')
da = xr.open_mfdataset(ifiles)['pvi'].assign_coords(time=time_ntc).sel(ens=ens)
pvi = da

#nseed
#ifile = 'tc_counts.seed.6ens.1871-2100.minohour12.nc'
ifile = glob.glob(f'tc_counts.seed*.minohour12.nc')[0]
ds = xr.open_dataset(ifile).rename(en='ens')
ds_seed = ds
nseed = ds[basin].sel(ens=ens)
#prob
prob = ntc/nseed
prob = prob.where(np.isfinite(prob), other=np.nan).interpolate_na(dim='time', fill_value='extrapolate').clip(max=1)


if __name__ == '__main__':
    from wyconfig import * #my plot settings
    import xfilter
    iframe = get_kws_from_argv('iframe', None)
    if iframe is not None: iframe = int(iframe)
    SI = True if 'SI' in sys.argv else False #for supported info
    fig,ax = plt.subplots()
    year_start,year_end = 1871,2100

    #ntc
    label = 'NTC'
    da = ntc.groupby('time.year').sum('time')
    da = da.mean('ens')
    da = da.pipe(norm).pipe(lowpass).pipe(units)
    if iframe is None or iframe >= 2:
        da.plot(label=label, ax=ax, color='k', lw=3)

    #nseed
    label = 'Nseed'
    da = nseed.groupby('time.month') * prob.sel(time=timeslice_ref).groupby('time.month').mean('time')
    da = da.groupby('time.year').sum('time')
    da = da.mean('ens')
    da = da.pipe(norm).pipe(lowpass).pipe(units)
    if iframe is None or iframe >= 3:
        da.plot(label=label, ax=ax, color='C0', ls='--')

    #spi
    label = 'SPI'
    da = spi.groupby('time.month') * prob.sel(time=timeslice_ref).groupby('time.month').mean('time')
    da = da.groupby('time.year').sum('time')
    da = da.mean('ens')
    da = da.pipe(norm).pipe(lowpass).pipe(units)
    if iframe is None or iframe >= 3:
        da.plot(label=label, ax=ax, color='C0', ls=':')

    #spi * prob
    label = 'SPIxProb'
    da = spi * prob
    da = da.groupby('time.year').sum('time')
    da = da.mean('ens')
    da = da.pipe(norm).pipe(lowpass).pipe(units)
    if iframe is None or iframe >= 4:
        da.plot(label=label, ax=ax, color='C0')

    #prob
    label = 'Prob'
    da = prob.groupby('time.month') * nseed.sel(time=timeslice_ref).groupby('time.month').mean('time')
    da = da.groupby('time.year').sum('time')
    da = da.mean('ens')
    da = da.pipe(norm).pipe(lowpass).pipe(units)
    if iframe is None or iframe >= 3:
        da.plot(label=label, ax=ax, color='C1', ls='--')

    #pvi
    label = 'P(VI)'
    da = pvi.groupby('time.month') * nseed.sel(time=timeslice_ref).groupby('time.month').mean('time')
    da = da.groupby('time.year').mean('time')
    da = da.mean('ens')
    da = da.pipe(norm).pipe(lowpass).pipe(units)
    if iframe is None or iframe >= 4:
        da.plot(label=label, ax=ax, color='C1', ls=':')

    # pvi * nseed
    label = 'P(VI)xNseed'
    da = pvi * nseed
    da = da.groupby('time.year').sum('time')
    da = da.mean('ens')
    da = da.pipe(norm).pipe(lowpass).pipe(units)
    if iframe is None or iframe >= 4:
        da.plot(label=label, ax=ax, color='C1')

    #spi * pvi
    label = 'SPIxP(VI)'
    da = spi * pvi
    da = da.groupby('time.year').sum('time')
    da = da.mean('ens')
    da = da.pipe(norm).pipe(lowpass).pipe(units)
    if iframe is None or iframe >= 4:
        da.plot(label=label, ax=ax, color='C2')

    ax.axhline(0, color='gray', ls='--')
    #ax.legend(loc='lower left', ncol=2)
    ax.legend(ncol=2)
    if len(ens)>1:
        title = f'{basin} TC, {n_window}-year-lowpass, {model}, {len(ens)}ens'
    else:
        title = f'{basin} TC, {n_window}-year-lowpass, {model}, ens{ens[0]}'
    #if minWindmax == 33: title = title.replace('TC', 'HU')
    ax.set_title(title)
    ax.set_xlim(year_start, year_end)
    ax.set_ylim(-40,60)

    #savefig
    if 'savefig' in sys.argv or 's' in sys.argv:
        figname = __file__.replace('.py', f'__{model}_{n_window}lowpass_{year_start:04d}-{year_end:04d}.png')
        if len(ens)>1:
            figname = figname.replace('.png', f'_{len(ens)}ens.png')
        else:
            figname = figname.replace('.png', f'_ens{ens[0]}.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 SI: 
            figname = figname.replace('.png', '_SI.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()