#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Wed Jun 30 17:02:45 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 scipy.stats import genextreme as gev, rankdata
from scipy.optimize import curve_fit, least_squares, minimize
from numpy import exp
#from numba import vectorize
import xfilter
lowpass = lambda x: x.filter.lowpass(1/5, dim='year', padtype='even')
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
ifile = '../misc/itxx_era5_index.nc'
da = xr.open_dataarray(ifile, decode_times=False).rename(time='year').assign_coords(year=range(1950,2022))
ts = da
ifile = '../misc/igiss_al_gl_m.nc'
da = xr.open_dataarray(ifile, decode_times=False)
da = da.assign_coords(time=pd.date_range('1880-01', freq='MS', periods=da.size))
da = da.groupby('time.year').mean('time').pipe(lowpass).sel(year=slice(1950,2021))
gmst = da
ds = xr.Dataset(dict(eindex=ts, gmst=gmst))

"""
y = da.sortby(da)
ep = 1 - (y.rank('year')-1)/y.size
rp = 1/ep
#fit by by scipy.stats.genextreme
xi, mu, sigma = gev.fit(da)
print(f'{xi = }; {mu = }; {sigma = }')
#fit by scipy.optimize.curve_fit
@vectorize(["float64(float64, float64, float64, float64)"], nopython=True)
def gev_cdf(x, mu, sigma, xi):
    if xi == 0:
        return exp( -exp( -(x-mu)/sigma ) )
    else: #xi != 0
        s = (x - mu)/sigma
        if xi*s > -1:
            return exp(-(1 + xi*(x - mu)/sigma)**(-1/xi)) 
        else: #xi*s <= -1
            if xi > 0:
                return 0.0
            else: #xi <0
                return 1.0
"""
def gev_cdf(x, mu, sigma, xi):
    #see https://en.wikipedia.org/wiki/Generalized_extreme_value_distribution
    if xi == 0:
        return exp( -exp( -(x-mu)/sigma ) )
    else: #xi != 0
        s = (x - mu)/sigma
        #return np.where(xi*s>-1, exp(-(1 + xi*(x - mu)/sigma)**(-1/xi)), (x*0 + 1 - np.sign(xi))/2)
        return (xi*s>-1)*exp(-(1 + xi*(x - mu)/sigma)**(-1/xi)) + (xi*s<=-1)*(xi<0)
def func_cost(pp, xx, tt):
    mu0, sigma, xi, alpha = pp
    xx_ = xx - alpha*tt
    cdf_emperical = rankdata(xx_)/xx_.size
    return gev_cdf(xx_, mu0, sigma, xi) - cdf_emperical
def minusLogLikelihood(params, xx, tt):
    mu, sigma, xi, alpha = params
    s = (xx - mu - alpha*tt)/sigma
    if xi == 0:
        return xx.size*log(sigma) + np.sum(s + exp(-s))
    else:
        return -np.sum( log(xi*s>-1) ) + xx.size*log(sigma) + np.sum( (1+1/xi)*log(1+xi*s) ) + np.sum( (1+xi*s)**(-1/xi) )
def gev_return_period(x, mu, sigma, xi):
    return 1.0/( 1 - gev_cdf(x, mu, sigma, xi) )
#r = least_squares(func_cost, x0=(20, 2, -0.2, 2), args=(ts.sortby(ts).values, gmst.sortby(ts).values), bounds=([0, 0, -0.4, 0], [40, 10, 0.4, 10])) 
x0 = [ts.mean().item(), ts.std().item(), -0.3, 2]
method = 'Nelder-Mead'
bounds = ( (None,None),  (0,None),  (None, None), (None,None))
res = minimize(minusLogLikelihood, x0, args=(ts.values[:-1], gmst.values[:-1]), method=method, bounds=bounds)#, options={'adaptive': True, 'maxfev': 1000} )
print(f'{res = }')
sys.exit()
    
xdata = y
ydata = y.rank('year')/y.size
popt, pcov = curve_fit(gev_cdf, xdata.values, ydata.values, bounds=([-np.inf, 0, -0.2],[np.inf, np.inf, 0.2]))
print(f'{popt = }')

xdata = y
ydata = 1/(1- (y.rank('year')-1)/y.size)
popt1, pcov = curve_fit(gev_return_period, xdata.values, ydata.values, bounds=([-np.inf, 0, -1/2],[np.inf, np.inf, 1/2]))
print(f'{popt1 = }')
 
 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    fig,axes = plt.subplots(2, 1, figsize=(6,5))
    y_fit = np.arange(21, 37)

    ax = axes[0]
    ax.plot(rp, y, ls='none', marker='o', fillstyle='none')

    ep_fit = 1 - gev.cdf(y_fit, xi, mu, sigma)
    rp_fit = 1/ep_fit
    ax.plot(rp_fit, y_fit)

    ep_fit = 1 - gev_cdf(y_fit, *popt)
    rp_fit = 1/ep_fit
    ax.plot(rp_fit, y_fit, ls='--')
    
    ax.set_xscale('log')
    ax.set_xlabel('return period [yr]')
    ax.set_ylabel('TXx [$^\circ$C]')

    ax = axes[1]
    ax.plot(y, ep, ls='none', marker='o', fillstyle='none')
    ep_fit = 1 - gev.cdf(y_fit, xi, mu, sigma)
    ax.plot(y_fit, ep_fit)

    ep_fit = 1 - gev_cdf(y_fit, *popt)
    ax.plot(y_fit, ep_fit, ls='--')

    ax.set_xlabel('TXx [$^\circ$C]')
    ax.set_ylabel('exceedance probability')


    #savefig
    if 'savefig' in sys.argv:
        figname = __file__.replace('.py', f'.png')
        wysavefig(figname)
    tt.check(f'**Done**')
    plt.show()