#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Tue May 17 16:47:55 EDT 2022
if __name__ == '__main__':
    import sys
    from misc.timer import Timer
    tt = Timer('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
import xlinregress
from scipy.stats import norm
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
ifile = 'wwa_t_ref_max_FLOR_HistRCP45_tigercpu_intelmpi_18_576PE_5ens_1860-2100_SouthAsia_index.nc'
ifile_cv = 'wwa_gmst_smoothed_FLOR_HistRCP45_tigercpu_intelmpi_18_576PE_5ens_1860-2100.nc'
da = xr.open_dataarray(ifile)
dacv = xr.open_dataarray(ifile_cv)
years_fit = None
#years_fit = slice('1950', '2022')
#years_fit = slice('1860', '2022')
if years_fit is None:
    years = da.time.dt.year.values
    year0, year1 = years[0], years[-1]
    years_fit = slice(f'{year0:04d}', f'{year1:04d}')
year_now = 2022
cv_now = dacv.sel(time=f'{year_now:04d}').mean('time').mean('ens').item()
cv_cold = cv_now - 1.2
cv_warm = cv_now + 0.8
rp_now_best = 100 #return period of the extreme event at current climate
low_extreme = False and True
print(f'{low_extreme = }')
print(f'{years_fit = }')

#fit
fmt = '.5g'
da_fit = da.sel(time=years_fit).stack(s=da.dims)
dacv_fit = dacv.sel(time=years_fit).stack(s=dacv.dims)
rg = da_fit.linregress.on(dacv_fit, dim='s') 
mu0_best = rg.intercept.item()
alpha_best = rg.slope.item()
sigma_best = (da_fit - rg.predicted).std().item()
"""
print()
print('**best fit**')
print(f'{mu0_best = :{fmt}}')
print(f'{sigma_best = :{fmt}}')
print(f'{alpha_best = :{fmt}}')
"""
#mu_best in cold(-1.2K), current and warm(0.8K) climate
mu_cold_best = mu0_best + alpha_best*cv_cold
mu_now_best = mu0_best + alpha_best*cv_now
mu_warm_best = mu0_best + alpha_best*cv_warm
"""
print(f'{mu_cold_best = :{fmt}}')
print(f'{mu_now_best = :{fmt}}')
print(f'{mu_warm_best = :{fmt}}')
"""

#bootstrap fit
nmc = 1000
mc_seed = 0
ci = 95 #%
q = [1/2-ci/100/2, 1/2+ci/100/2]
rng = np.random.default_rng(mc_seed)
imc = rng.choice(da_fit.size, size=(nmc, da_fit.size))
da_fit_mc = xr.DataArray(da_fit.values[imc], dims=('mc', 's'))
dacv_fit_mc = xr.DataArray(dacv_fit.values[imc], dims=('mc', 's'))
rg_mc = da_fit_mc.linregress.on(dacv_fit_mc, dim='s')
mu0 = rg_mc.intercept
alpha = rg_mc.slope
sigma = (da_fit_mc - rg_mc.predicted).std('s')
mu0_ci = mu0.quantile(q, dim='mc').values
alpha_ci = alpha.quantile(q, dim='mc').values
sigma_ci = sigma.quantile(q, dim='mc').values
print()
print(f'**bootstrap fit resuts**: {mc_seed = }; {nmc = }')
n = 45
print(f'mu0 best estimate and bootstrap {ci}% CI:'.rjust(n), f'{mu0_best:{fmt}} and {mu0_ci[0]:{fmt}} ... {mu0_ci[-1]:{fmt}}')
print(f'sigma best estimate and bootstrap {ci}% CI:'.rjust(n), f'{sigma_best:{fmt}} and {sigma_ci[0]:{fmt}} ... {sigma_ci[-1]:{fmt}}')
print(f'alpha best estimate and bootstrap {ci}% CI:'.rjust(n), f'{alpha_best:{fmt}} and {alpha_ci[0]:{fmt}} ... {alpha_ci[-1]:{fmt}}')
print()
#mu in cold(-1.2K), current and warm(0.8K) climate
mu_cold = mu0 + alpha*cv_cold
mu_cold_ci = mu_cold.quantile(q, dim='mc').values
mu_now = mu0 + alpha*cv_now
mu_now_ci = mu_now.quantile(q, dim='mc').values
mu_warm = mu0 + alpha*cv_warm
mu_warm_ci = mu_warm.quantile(q, dim='mc').values
print(f'mu_cold best estimate and bootstrap {ci}% CI:'.rjust(n), f'{mu_cold_best:{fmt}} and {mu_cold_ci[0]:{fmt}} ... {mu_cold_ci[-1]:{fmt}}')
print(f'mu_now best estimate and bootstrap {ci}% CI:'.rjust(n), f'{mu_now_best:{fmt}} and {mu_now_ci[0]:{fmt}} ... {mu_now_ci[-1]:{fmt}}')
print(f'mu_warm best estimate and bootstrap {ci}% CI:'.rjust(n), f'{mu_warm_best:{fmt}} and {mu_warm_ci[0]:{fmt}} ... {mu_warm_ci[-1]:{fmt}}')

#return level of the event in model
if low_extreme:
    return_level = norm.ppf(1/rp_now_best, loc=mu_now_best, scale=sigma_best)
else:
    return_level = norm.ppf(1-1/rp_now_best, loc=mu_now_best, scale=sigma_best)
print()
print(f'{return_level = :{fmt}} for the extreme event in model')
#return period best estimate
if low_extreme:
    prob_cold_best = norm.cdf(return_level, loc=mu_cold_best, scale=sigma_best)
    prob_warm_best = norm.cdf(return_level, loc=mu_warm_best, scale=sigma_best)
else:
    prob_cold_best = 1 - norm.cdf(return_level, loc=mu_cold_best, scale=sigma_best)
    prob_warm_best = 1 - norm.cdf(return_level, loc=mu_warm_best, scale=sigma_best)
rp_cold_best = 1/prob_cold_best
rp_warm_best = 1/prob_warm_best
prob_now_best = 1/rp_now_best
"""
print()
print('**return period best estimate**:')
print(f'{rp_cold_best = :{fmt}}')
print(f'{rp_now_best = :{fmt}}')
print(f'{rp_warm_best = :{fmt}}')
"""

#return period and prob mc
if low_extreme:
    prob_cold = norm.cdf(return_level, loc=mu_cold, scale=sigma)
    prob_now = norm.cdf(return_level, loc=mu_now, scale=sigma)
    prob_warm = norm.cdf(return_level, loc=mu_warm, scale=sigma)
else:
    prob_cold = 1 - norm.cdf(return_level, loc=mu_cold, scale=sigma)
    prob_now = 1 - norm.cdf(return_level, loc=mu_now, scale=sigma)
    prob_warm = 1 - norm.cdf(return_level, loc=mu_warm, scale=sigma)
rp_cold = 1/prob_cold
rp_now = 1/prob_now
rp_warm = 1/prob_warm
#return period and prob CI
prob_cold_ci = np.quantile(prob_cold, q)
rp_cold_ci = np.quantile(rp_cold, q)
prob_now_ci = np.quantile(prob_now, q)
rp_now_ci = np.quantile(rp_now, q)
prob_warm_ci = np.quantile(prob_warm, q)
rp_warm_ci = np.quantile(rp_warm, q)
print()
print('**probability in cold, current and warm climate**')
print(f'prob_cold best estimate and bootstrap {ci}% CI:'.rjust(n), f'{prob_cold_best:{fmt}} and {prob_cold_ci[0]:{fmt}} ... {prob_cold_ci[-1]:{fmt}}')
print(f'prob_now best estimate and bootstrap {ci}% CI:'.rjust(n), f'{prob_now_best:{fmt}} and {prob_now_ci[0]:{fmt}} ... {prob_now_ci[-1]:{fmt}}')
print(f'prob_warm best estimate and bootstrap {ci}% CI:'.rjust(n), f'{prob_warm_best:{fmt}} and {prob_warm_ci[0]:{fmt}} ... {prob_warm_ci[-1]:{fmt}}')
print()
print('**return period in cold, current and warm climate**')
print(f'rp_cold best estimate and bootstrap {ci}% CI:'.rjust(n), f'{rp_cold_best:{fmt}} and {rp_cold_ci[0]:{fmt}} ... {rp_cold_ci[-1]:{fmt}}')
print(f'rp_now best estimate and bootstrap {ci}% CI:'.rjust(n), f'{rp_now_best:{fmt}} and {rp_now_ci[0]:{fmt}} ... {rp_now_ci[-1]:{fmt}}')
print(f'rp_warm best estimate and bootstrap {ci}% CI:'.rjust(n), f'{rp_warm_best:{fmt}} and {rp_warm_ci[0]:{fmt}} ... {rp_warm_ci[-1]:{fmt}}')

#prob ratio
pr_now_to_cold_best = prob_now_best/prob_cold_best
pr_warm_to_now_best = prob_warm_best/prob_now_best
print()
"""
print(f'{pr_now_to_cold_best = :{fmt}}')
print(f'{pr_warm_to_now_best = :{fmt}}')
"""
#CI
pr_now_to_cold = prob_now/prob_cold
pr_now_to_cold_ci = np.quantile(pr_now_to_cold, q)
pr_warm_to_now = prob_warm/prob_now
pr_warm_to_now_ci = np.quantile(pr_warm_to_now, q)
print('**probability ratio**')
print(f'pr_now_to_cold best estimate and bootstrap {ci}% CI:'.rjust(n), f'{pr_now_to_cold_best:{fmt}} and {pr_now_to_cold_ci[0]:{fmt}} ... {pr_now_to_cold_ci[-1]:{fmt}}')
print(f'pr_warm_to_now best estimate and bootstrap {ci}% CI:'.rjust(n), f'{pr_warm_to_now_best:{fmt}} and {pr_warm_to_now_ci[0]:{fmt}} ... {pr_warm_to_now_ci[-1]:{fmt}}')
 
#intensity change
cold_to_now_best = alpha_best * (cv_now - cv_cold)
now_to_warm_best = alpha_best * (cv_warm - cv_now)
print()
"""
print(f'{cold_to_now_best = :{fmt}}')
print(f'{now_to_warm_best = :{fmt}}')
"""
#CI
cold_to_now = alpha * (cv_now - cv_cold)
now_to_warm = alpha * (cv_warm - cv_now)
cold_to_now_ci = np.quantile(cold_to_now, q)
now_to_warm_ci = np.quantile(now_to_warm, q)
print('**intensity change**')
print(f'cold_to_now best estimate and bootstrap {ci}% CI:'.rjust(n), f'{cold_to_now_best:{fmt}} and {cold_to_now_ci[0]:{fmt}} ... {cold_to_now_ci[-1]:{fmt}}')
print(f'now_to_warm best estimate and bootstrap {ci}% CI:'.rjust(n), f'{now_to_warm_best:{fmt}} and {now_to_warm_ci[0]:{fmt}} ... {now_to_warm_ci[-1]:{fmt}}')


if __name__ == '__main__':
    from wyconfig import * #my plot settings
    from wyextreme.shared import plot_smp_return_period
    plt.close('all')
    def plot_return_period(loc, scale, loc_mc=None, scale_mc=None, ci=95, lower=1/.99, upper=10000, ax=None, low_extreme=False, label=None, **kws):
        #lower, upper = 1/.99, 1000
        #loc = mu_cold_best
        #scale = sigma_best
        xx = np.linspace( norm.ppf(1-1/lower, loc=loc, scale=scale),
            norm.ppf(1-1/upper, loc=loc, scale=scale), 100)
        if low_extreme:
            prob = norm.cdf(xx, loc=loc, scale=scale)
        else:
            prob = 1 - norm.cdf(xx, loc=loc, scale=scale)
        rp = 1/prob
        ax.plot(rp, xx, label=label, **kws)

        if loc_mc is not None:
            ns = 100
            rp = np.logspace(np.log10(1.1), np.log10(upper), ns)
            if low_extreme:
                xx = norm.ppf(1/rp, loc=np.array(loc_mc).reshape(loc_mc.size,1), scale=np.array(scale_mc).reshape(scale_mc.size,1))
            else:
                xx = norm.ppf(1-1/rp, loc=np.array(loc_mc).reshape(loc_mc.size,1), scale=np.array(scale_mc).reshape(scale_mc.size,1))
            xx_ci = np.quantile(xx, q, axis=0)
            ax.plot(rp, xx_ci[0], lw=0.5, **kws)
            ax.plot(rp, xx_ci[-1], lw=0.5, **kws)
    
    fig, ax = plt.subplots()
    plot_smp_return_period(da_fit + (cv_cold-dacv_fit)*alpha_best, ax=ax, color='C0', low_extreme=low_extreme)
    plot_smp_return_period(da_fit + (cv_now-dacv_fit)*alpha_best, ax=ax, color='C1', low_extreme=low_extreme)
    plot_smp_return_period(da_fit + (cv_warm-dacv_fit)*alpha_best, ax=ax, color='C4', low_extreme=low_extreme)
    ax.axhline(return_level, color='gray', ls='--')
    
    plot_return_period(loc=mu_cold_best, scale=sigma_best, ax=ax, color='C0', loc_mc=mu_cold, scale_mc=sigma, low_extreme=low_extreme, label='$-1.2$')
    plot_return_period(loc=mu_now_best, scale=sigma_best, ax=ax, color='C1', loc_mc=mu_now, scale_mc=sigma, low_extreme=low_extreme, label=f'{year_now}')
    plot_return_period(loc=mu_warm_best, scale=sigma_best, ax=ax, color='C4', loc_mc=mu_warm, scale_mc=sigma, low_extreme=low_extreme, label='$+0.8$')

    ax.legend()
    ax.set_xlabel('return period [yr]')
    ax.set_ylabel(f'return level [{da.attrs["units"]}]')
    ax.set_title(f'{years_fit.start}-{years_fit.stop}')

    #savefig
    if len(sys.argv)>1 and 'savefig' in sys.argv[1:]:
        figname = __file__.replace('.py', f'.png')
        if 'overwritefig' in sys.argv[1:]:
            wysavefig(figname, overwritefig=True)
        else:
            wysavefig(figname)

    #plot scatter plot
    fig, ax = plt.subplots()
    ax.scatter(dacv_fit, da_fit, color='none', edgecolor='gray', alpha=0.5, s=10)
    ax.axline((cv_now, mu_now_best), slope=alpha_best, color='k')
    if low_extreme:
        x6 = norm.ppf(1/6, loc=mu_now_best, scale=sigma_best)
        x40 = norm.ppf(1/40, loc=mu_now_best, scale=sigma_best)
    else:
        x6 = norm.ppf(1-1/6, loc=mu_now_best, scale=sigma_best)
        x40 = norm.ppf(1-1/40, loc=mu_now_best, scale=sigma_best)
    ax.axline((cv_now, x6), slope=alpha_best, lw=0.5, color='k')
    ax.axline((cv_now, x40), slope=alpha_best, lw=0.5, color='k')

    capsize = 2
    ax.errorbar(cv_cold, mu_cold_best, np.abs(mu_cold_ci-mu_cold_best).reshape(2,1), capsize=capsize, color='C0')
    ax.errorbar(cv_now, mu_now_best, np.abs(mu_now_ci-mu_now_best).reshape(2,1), capsize=capsize, color='C1')
    ax.errorbar(cv_warm, mu_warm_best, np.abs(mu_warm_ci-mu_warm_best).reshape(2,1), capsize=capsize, color='C4')

    ax.set_xlabel('GMST [K]')
    ax.set_ylabel(f'{da.attrs["units"]}')

    tt.check(f'**Done**')
    print()
    if 'notshowfig' in sys.argv:
        pass
    else:
        plt.show()