#!/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 {__file__}')
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
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
#years = slice(1950, None)
n_ens = 10
print(f'{n_ens = }')
years = slice(1901, 2014)
years_fit = slice(1901, 2014)
#years_fit = slice(1901, 2014)
#years_fit = slice(1980, 2020)
#years_fit = slice(1880, 2020)
#non-urban
ifile = './TXx_AM4urban_wasteCool_0urban_amip_5ens_1870-2014.nc'
txx_n = xr.open_dataarray(ifile).sel(year=years).isel(en=slice(None, n_ens)).stack(s=['en', 'year'])
txx_n_fit = xr.open_dataarray(ifile).sel(year=years_fit).isel(en=slice(None, n_ens)).stack(s=['en', 'year'])
ifile = './GMST_AM4urban_wasteCool_0urban_amip_5ens_1870-2014.nc'
#gmst_n = xr.open_dataarray(ifile).rolling(year=5, center=True, min_periods=1).mean().sel(year=years).stack(s=['en', 'year'])
gmst_n = xr.open_dataarray(ifile).rolling(year=5, min_periods=1).mean().sel(year=years).isel(en=slice(None, n_ens)).stack(s=['en', 'year']) #center=False
gmst_n_fit = xr.open_dataarray(ifile).rolling(year=5, min_periods=1).mean().sel(year=years_fit).isel(en=slice(None, n_ens)).stack(s=['en', 'year']) #center=False
#gmst_n = xr.open_dataarray(ifile).rolling(year=5, min_periods=1).mean().sel(year=years) #center=False
#ens = gmst_n.en
#gmst_n = xr.concat([gmst_n.mean('en'),]*ens.size, dim=pd.Index(ens.values, name='en')).stack(s=['en', 'year'])
#urban
ifile = './TXx_AM4urban_wasteCool_amip_5ens_1870-2014.nc'
txx_u = xr.open_dataarray(ifile).sel(year=years).isel(en=slice(None, n_ens)).stack(s=['en', 'year'])
txx_u_fit = xr.open_dataarray(ifile).sel(year=years_fit).isel(en=slice(None, n_ens)).stack(s=['en', 'year'])
ifile = './GMST_AM4urban_wasteCool_amip_5ens_1870-2014.nc'
#gmst_u = xr.open_dataarray(ifile).rolling(year=5, center=True, min_periods=1).mean().sel(year=years).stack(s=['en', 'year'])
gmst_u = xr.open_dataarray(ifile).rolling(year=5, min_periods=1).mean().sel(year=years).isel(en=slice(None, n_ens)).stack(s=['en', 'year']) #center=False
gmst_u_fit = xr.open_dataarray(ifile).rolling(year=5, min_periods=1).mean().sel(year=years_fit).isel(en=slice(None, n_ens)).stack(s=['en', 'year']) #center=False
#gmst_u = xr.open_dataarray(ifile).rolling(year=5, min_periods=1).mean().sel(year=years) #center=False
#ens = gmst_u.en
#gmst_u = xr.concat([gmst_u.mean('en'),]*ens.size, dim=pd.Index(ens.values, name='en')).stack(s=['en', 'year'])
"""
#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
    """
    print('no urban')
    fit_all(txx_n, gmst_n)

    print('urban')
    fit_all(txx_u, gmst_u)
    """
    fig, ax = plt.subplots()
    upper_rp = 1e6
    nmc = 200
    fit_kws = {'xi_bounds': (-0.4, 0.4)}
    print('no urban')
    year = gmst_n.year[-1].item()
    #ds_n = plot_fit_bootstrap(txx_n, gmst_n, ('no-urban -1.2', gmst_n[-1].item()-1.2), ax=ax, color='C0', upper_rp=upper_rp, 
    #    nmc=nmc, fit_kws=fit_kws)
    ds_n = plot_fit_bootstrap(txx_n_fit, gmst_n_fit, (f'no-urban {year}', gmst_n[-1].item()), bsfit=None, ax=ax, color='C2', upper_rp=upper_rp, 
        nmc=nmc, fit_kws=fit_kws)
    print('urban')
    year = gmst_u.year[-1].item()
    ds_u = plot_fit_bootstrap(txx_u_fit, gmst_u_fit, ('urban -1.2', gmst_u[-1].item()-1.2), ax=ax, color='C0', upper_rp=upper_rp, 
        nmc=nmc, fit_kws=fit_kws)
    ds_u = plot_fit_bootstrap(txx_u_fit, gmst_u_fit, (f'urban {year}', gmst_u[-1].item()), bsfit=ds_u, ax=ax, color='C3', upper_rp=upper_rp, 
    nmc=nmc, fit_kws=fit_kws)

    rlevel = gev_return_period_inverse(1000,  ds_u.mu0_best + gmst_u[-1]*ds_u.alpha_best, ds_u.sigma_best, ds_u.xi_best) #return level for 1000 return years
    ax.axhline(rlevel, color='gray', ls='--')
    ax.legend()
    ax.set_xlabel('return period [yr]')
    ax.set_ylabel('PNW TXx index [$^\circ$C]')

    print()
    print(f'{years_fit = }')
    #probability ratio of urban over non-urban
    rp_n = gev_return_period(rlevel, (ds_n.mu0_best + gmst_n[-1]*ds_n.alpha_best), ds_n.sigma_best, ds_n.xi_best).item()
    prob_ratio_best = rp_n/1000
    print(f'probability ratio of urban over non-urban: {prob_ratio_best:.4g}') 
    rp_n = [gev_return_period(rlevel, (ds_n.mu0 + gmst_n[-1]*ds_n.alpha)[ii], ds_n.sigma[ii], ds_n.xi[ii]).item() for ii in range(ds_n.xi.size)]
    rp_n = xr.DataArray(rp_n, dims='mc').fillna(np.inf)
    prob_ratio = (rp_n/1000).quantile(.5).fillna(np.inf)
    print(f'probability ratio of urban over non-urban median: {prob_ratio.item():.4g}')
    prob_ratio = (rp_n/1000).quantile([.25, .75]).fillna(np.inf)
    print(f'probability ratio of urban over non-urban 50% CI: {prob_ratio[0].item():.4g}, {prob_ratio[1].item():.4g}')
    prob_ratio = (rp_n/1000).quantile([.05, .95]).fillna(np.inf)
    print(f'probability ratio of urban over non-urban 90% CI: {prob_ratio[0].item():.4g}, {prob_ratio[1].item():.4g}')
    prob_ratio = (rp_n/1000).quantile([.025, .975]).fillna(np.inf)
    print(f'probability ratio of urban over non-urban 95% CI: {prob_ratio[0].item():.4g}, {prob_ratio[1].item():.4g}')
    
    #probability ratio of current climate over PI climamte 
    print()
    rp_pi = gev_return_period(rlevel, (ds_u.mu0_best + (gmst_u[-1]-1.2)*ds_u.alpha_best), ds_u.sigma_best, ds_u.xi_best).item()
    prob_ratio_best = rp_pi/1000
    prob_ratio_best = np.inf if np.isnan(prob_ratio_best) else prob_ratio_best
    year = gmst_u.year[-1].item()
    print(f'probability ratio of {year} over PI climate: {prob_ratio_best:.4g}') 
    rp_pi = [gev_return_period(rlevel, (ds_u.mu0 + (gmst_u[-1]-1.2)*ds_u.alpha)[ii], ds_u.sigma[ii], ds_u.xi[ii]).item() for ii in range(ds_u.xi.size)]
    rp_pi = xr.DataArray(rp_pi, dims='mc').fillna(np.inf)
    prob_ratio = (rp_pi/1000).quantile(.5).fillna(np.inf)
    print(f'probability ratio of {year} over PI climate median: {prob_ratio.item():.4g}')
    prob_ratio = (rp_pi/1000).quantile([.25, .75]).fillna(np.inf)
    print(f'probability ratio of {year} over PI climate 50% CI: {prob_ratio[0].item():.4g}, {prob_ratio[1].item():.4g}')
    prob_ratio = (rp_pi/1000).quantile([.05, .95]).fillna(np.inf)
    print(f'probability ratio of {year} over PI climate 90% CI: {prob_ratio[0].item():.4g}, {prob_ratio[1].item():.4g}')
    prob_ratio = (rp_pi/1000).quantile([.025, .975]).fillna(np.inf)
    print(f'probability ratio of {year} over PI climate 95% CI: {prob_ratio[0].item():.4g}, {prob_ratio[1].item():.4g}')
    
    #savefig
    if 'savefig' in sys.argv:
        figname = __file__.replace('.py', f'_return_period_{n_ens}ens.png')
        if years_fit.start>1901:
            figname = figname.replace('.png', f'_{years_fit.start}.png')
        if years_fit.stop<2020:
            figname = figname.replace('.png', f'_{years_fit.stop}.png')
        wysavefig(figname)

    fig, ax = plt.subplots()
    plot_covariate(txx_u_fit, gmst_u_fit, ax=ax, color='k')
    plot_mu_ci(txx_u_fit, gmst_u_fit, gmst_u[-1]-1.2, bsfit=ds_u)
    plot_mu_ci(txx_u_fit, gmst_u_fit, gmst_u[-1], bsfit=ds_u)
    #plot_covariate(txx_n, gmst_n, ax=ax, color='gray')
    ax.set_ylabel('PNW TXx index [$^\circ$C]')
    ax.set_xlabel('Global mean surface temperature [$^\circ$C]')

    #savefig
    if 'savefig' in sys.argv:
        figname = __file__.replace('.py', f'_covariate_{n_ens}ens.png')
        if years_fit.start>1901:
            figname = figname.replace('.png', f'_{years_fit.start}.png')
        if years_fit.stop<2020:
            figname = figname.replace('.png', f'_{years_fit.stop}.png')
        wysavefig(figname)
    tt.check(f'**Done**')
    plt.show()