#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Fri Jul 16 11:49:12 EDT 2021
if __name__ == '__main__':
    from misc.timer import Timer
    tt = Timer(f'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
#from wyextreme.gev_shift import plot_fit_bootstrap, fit_all, gev_return_period, gev_return_period_inverse, plot_covariate, plot_mu_ci
from wyextreme.gev_scale import plot_fit_bootstrap, fit_all, gev_return_period, gev_return_period_inverse, plot_covariate, plot_mu_ci
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
year_event = 2022
ndays = 15 if '15' in sys.argv else 7 #7(or 15)-day running average
return_period = 500 if ndays==7 else 1000
nmc = 1000 if '1000' in sys.argv else 100# [100, 500, 1000][0] #MC sample size
upper_rp = 1e4
fit_kws = None # {'xi_bounds': (-0.4, 0.4)}

#func_anomaly = lambda x: x - x.mean('year')
func_anomaly = lambda x: x - x.sel(year=slice(1951,1980)).mean('year')
years = slice(1861, 2100)
years_fit = slice(1861, 2100)
#years_fit = slice(1960, 2022)
ifile = f'../wwa_precip_FLOR_HistRCP45_tigercpu_intelmpi_18_576PE_5ens_1860-2100_Pernambuco_{ndays}days_index.nc'
da = xr.open_dataarray(ifile).groupby('time.year').mean('time').sel(year=years).stack(s=['ens', 'year'])
da_fit = xr.open_dataarray(ifile).groupby('time.year').mean('time').sel(year=years_fit).stack(s=['ens', 'year'])
#ifile = '../GMST_AM4urban_wasteCool_amip_5ens_1870-2020.nc'
#ifile = '../data_gmst_AM2p5C360_amipLongChan_10ens_1872-2020_yearlymeanJul-Jun.nc'
ifile = '../wwa_gmst_smoothed_FLOR_HistRCP45_tigercpu_intelmpi_18_576PE_5ens_1860-2100.nc'
gmst = xr.open_dataarray(ifile).groupby('time.year').mean('time').pipe(func_anomaly).rolling(year=5, min_periods=1).mean().sel(year=years).stack(s=['ens', 'year']) #center=False
gmst_fit = xr.open_dataarray(ifile).groupby('time.year').mean('time').pipe(func_anomaly).rolling(year=5, min_periods=1).mean().sel(year=years_fit).stack(s=['ens', 'year']) #center=False
#wwa format gmst already anomaly and smoothed
#gmst = xr.open_dataarray(ifile).groupby('time.year').mean('time').sel(year=years).stack(s=['ens', 'year']) #
#gmst_fit = xr.open_dataarray(ifile).groupby('time.year').mean('time').sel(year=years_fit).stack(s=['ens', 'year']) #
gmst_event = gmst.sel(year=year_event).mean('ens').item()
"""
#test
gmst_u = xr.open_dataarray(ifile).rolling(year=5, min_periods=1).mean().sel(year=years).stack(s=['en', 'year']) #center=False
xr.open_dataarray(ifile).rolling(year=5, min_periods=1).mean().sel(year=slice(1990,None)).mean('en').plot(marker='o', fillstyle='none', hue='en')
plt.grid('on')
plt.show()
sys.exit()
"""

 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    fig, ax = plt.subplots()
    #upper_rp = 1e6
    #fit plots
    ds = plot_fit_bootstrap(da_fit, gmst_fit, ('$-$1.2', gmst_event-1.2), ax=ax, color='C0', upper_rp=upper_rp, 
        nmc=nmc, fit_kws=fit_kws)
    ds = plot_fit_bootstrap(da_fit, gmst_fit, (f'{year_event}', gmst_event), bsfit=ds, ax=ax, color='C1', upper_rp=upper_rp, 
        nmc=nmc, fit_kws=fit_kws)
    
    mu0, sigma0, xi, alpha = ds.mu0_best, ds.sigma0_best, ds.xi_best, ds.alpha_best
    rlevel = gev_return_period_inverse(return_period,  mu0*np.exp(gmst_event*alpha/mu0), sigma0*np.exp(gmst_event*alpha/mu0), xi) #return level for 20 return years
    print(f'threshold of return period {return_period}: {rlevel}')
    ax.axhline(rlevel, color='gray', ls='--')
    ax.legend()
    ax.set_xlabel('return period [yr]')
    ax.set_ylabel(f'Pernambuco {ndays}-day precip index [mm/day]')

    print()
    print(f'{years_fit = }')

    #probability ratio of current climate over PI climamte 
    print()
    rp_pi = gev_return_period(rlevel, mu0*np.exp((gmst_event-1.2)*alpha/mu0), sigma0*np.exp((gmst_event-1.2)*alpha/mu0), xi).item()
    prob_ratio_best = rp_pi/return_period
    prob_ratio_best = np.inf if np.isnan(prob_ratio_best) else prob_ratio_best
    print(f'probability ratio of {year_event} over PI climate: {prob_ratio_best:.4g}') 
    rp_pi = [gev_return_period(rlevel, (ds.mu0*np.exp((gmst_event-1.2)*ds.alpha/ds.mu0))[ii], (ds.sigma0*np.exp((gmst_event-1.2)*ds.alpha/ds.mu0))[ii], ds.xi[ii]).item() 
        for ii in range(ds.xi.size)]
    rp_pi = xr.DataArray(rp_pi, dims='mc').fillna(np.inf)
    prob_ratio = (rp_pi/return_period).quantile(.5).fillna(np.inf)
    print(f'probability ratio of {year_event} over PI climate median: {prob_ratio.item():.4g}')
    prob_ratio = (rp_pi/return_period).quantile([.25, .75]).fillna(np.inf)
    print(f'probability ratio of {year_event} over PI climate 50% CI: {prob_ratio[0].item():.4g}, {prob_ratio[1].item():.4g}')
    prob_ratio = (rp_pi/return_period).quantile([.05, .95]).fillna(np.inf)
    print(f'probability ratio of {year_event} over PI climate 90% CI: {prob_ratio[0].item():.4g}, {prob_ratio[1].item():.4g}')
    prob_ratio = (rp_pi/return_period).quantile([.025, .975]).fillna(np.inf)
    print(f'probability ratio of {year_event} over PI climate 95% CI: {prob_ratio[0].item():.4g}, {prob_ratio[1].item():.4g}')
    
    #probability ratio of 2K warming to current climate
    print()
    rp_2k = gev_return_period(rlevel, mu0*np.exp((gmst_event+0.8)*alpha/mu0), sigma0*np.exp((gmst_event+0.8)*alpha/mu0), xi).item()
    prob_ratio_best = return_period/rp_2k
    prob_ratio_best = np.inf if np.isnan(prob_ratio_best) else prob_ratio_best
    print(f'probability ratio of 2K warming climate to {year_event}: {prob_ratio_best:.4g}') 
    rp_2k = [gev_return_period(rlevel, (ds.mu0*np.exp((gmst_event+0.8)*ds.alpha/ds.mu0))[ii], (ds.sigma0*np.exp((gmst_event+0.8)*ds.alpha/ds.mu0))[ii], ds.xi[ii]).item() 
        for ii in range(ds.xi.size)]
    rp_2k = xr.DataArray(rp_2k, dims='mc').fillna(np.inf)
    prob_ratio = (return_period/rp_2k).quantile(.5).fillna(np.inf)
    print(f'probability ratio of 2K warming to {year_event} climate median: {prob_ratio.item():.4g}')
    prob_ratio = (return_period/rp_2k).quantile([.25, .75]).fillna(np.inf)
    print(f'probability ratio of 2K warming to {year_event} climate 50% CI: {prob_ratio[0].item():.4g}, {prob_ratio[1].item():.4g}')
    prob_ratio = (return_period/rp_2k).quantile([.05, .95]).fillna(np.inf)
    print(f'probability ratio of 2K warming to {year_event} climate 90% CI: {prob_ratio[0].item():.4g}, {prob_ratio[1].item():.4g}')
    prob_ratio = (return_period/rp_2k).quantile([.025, .975]).fillna(np.inf)
    print(f'probability ratio of 2K warming to {year_event} climate 95% CI: {prob_ratio[0].item():.4g}, {prob_ratio[1].item():.4g}')
    
    print()
    print('intensity change from PI climate:') #units are %
    dI_best = 100*( np.exp(1.2*ds.alpha_best/ds.mu0_best) -1 ).item()
    dI_ci = 100*( np.exp(1.2*ds.alpha/ds.mu0) -1 ).quantile([0.025, 0.975], dim='mc')
    print(f'Intensity change and 95% CI:'.rjust(20), f'{dI_best:.4g}({dI_ci[0].item():.4g}, {dI_ci[1].item():.4g})')
    print()

    print('intensity change to 2K warming climate:') #units are %
    dI_best = 100*( np.exp(0.8*ds.alpha_best/ds.mu0_best) -1 ).item()
    dI_ci = 100*( np.exp(0.8*ds.alpha/ds.mu0) -1 ).quantile([0.025, 0.975], dim='mc')
    print(f'Intensity change and 95% CI:'.rjust(20), f'{dI_best:.4g}({dI_ci[0].item():.4g}, {dI_ci[1].item():.4g})')
    print()
    
    #savefig
    if 'savefig' in sys.argv:
        figname = __file__.replace('.py', f'_{ndays}days_{years_fit.start}-{years_fit.stop}.png')
        wysavefig(figname)

    fig, ax = plt.subplots()
    plot_covariate(da_fit, gmst_fit, ax=ax)
    print('PI climate')
    plot_mu_ci(da_fit, gmst_fit, gmst_event-1.2, bsfit=ds, color='C0')
    print()
    print('Current climate')
    plot_mu_ci(da_fit, gmst_fit, gmst_event, bsfit=ds, color='C1')
    print()
    print('2K warming climate')
    plot_mu_ci(da_fit, gmst_fit, gmst_event+0.8, bsfit=ds, color='C3')
    #plot_covariate(txx_n, gmst_n, ax=ax, color='gray')
    ax.set_ylabel(f'Pernambuco {ndays}-day precip index [mm/day]')
    ax.set_xlabel('Global mean surface temperature [$^\circ$C]')

    #savefig
    if 'savefig' in sys.argv:
        figname = __file__.replace('.py', f'_covariate_{ndays}days_{years_fit.start}-{years_fit.stop}.png')
        wysavefig(figname)
    tt.check(f'**Done**')
    plt.show()