#!/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
from scipy.optimize import curve_fit, minimize
from numpy import exp, log
from numba import vectorize
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
ifile = '../TXx_FLOR_CTL1990_0101-1000.nc'
da = xr.open_dataarray(ifile)
y = da.sortby(da)
ep = 1 - (y.rank('year')-1)/y.size
rp = 1/ep
#fit by by scipy.stats.genextreme
c, 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)
#pdf
def gev_pdf(x, mu, sigma, xi, dx=1e-10):
    return ( gev_cdf(x+dx, mu, sigma, xi) - gev_cdf(x-dx, mu, sigma, xi) )/(2*dx)
def gev_pdf_true(x, mu, sigma, xi):
    s = (x-mu)/sigma
    if xi == 0:
        tx = exp(-s)
    else:
        tx = (1+xi*s)**(-1/xi)
    return (xi*s>-1)*(1/sigma) * (tx**(1+xi)) * exp(-tx)
def minusLogLikelihood(params, xx):
    mu, sigma, xi = params
    s = (xx - mu)/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) )
        #return xx.size*log(sigma) + np.sum( (1+1/xi)*log(1+xi*s) ) + np.sum( (1+xi*s)**(-1/xi) )

x0 = [da.mean().item(), da.std().item(), 0]
method = ['Nelder-Mead',][0]#, 'Powell', 'SLSQP','BFGS', 'L-BFGS-B'][0]
bounds = ( (None,None),  (0,None),  (None, None))
res = minimize(minusLogLikelihood, x0, args=(da.values), method=method, bounds=bounds )
print(f'{res = }')
print(f'{x0 = }')
print(f'{res.x = }')
print(f'{mu = }; {sigma = }; {-c = }')
        
#r"""
x = np.linspace(-4,4,100)
#cdf
plt.figure()
plt.plot(x, gev_cdf(x, 0, 1, -1/2), label=r'$\xi$ = -1/2')
plt.plot(x, gev_cdf(x, 0, 1, 0), label=r'$\xi$ = 0')
plt.plot(x, gev_cdf(x, 0, 1, 1/2), label=r'$\xi$ = 1/2')
plt.legend()
#pdf
plt.figure()
plt.plot(x, gev_pdf(x, 0, 1, -1/2), label=r'$\xi$ = -1/2')
plt.plot(x, gev_pdf(x, 0, 1, 0), label=r'$\xi$ = 0')
plt.plot(x, gev_pdf(x, 0, 1, 1/2), label=r'$\xi$ = 1/2')
plt.plot(x, gev_pdf_true(x, 0, 1, -1/2), label=r'$\xi$ = -1/2', ls='--')
plt.plot(x, gev_pdf_true(x, 0, 1, 0), label=r'$\xi$ = 0', ls='--')
plt.plot(x, gev_pdf_true(x, 0, 1, 1/2), label=r'$\xi$ = 1/2', ls='--')
plt.legend()
plt.show()
#"""

def gev_return_period(x, mu, sigma, xi):
    return 1.0/( 1 - gev_cdf(x, mu, sigma, xi) )


xdata = y
ydata = y.rank('year')/y.size
popt, pcov = curve_fit(gev_cdf, xdata.values, ydata.values, bounds=([-np.inf, 0, -0.5],[np.inf, np.inf, 0.5]))
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, c, mu, sigma)
    rp_fit = 1/ep_fit
    ax.plot(rp_fit, y_fit, label='scipy.stats.genextreme fit')

    ep_fit = 1 - gev_cdf(y_fit, *popt)
    rp_fit = 1/ep_fit
    ax.plot(rp_fit, y_fit, ls='--', label='cdf curve_fit')

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

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

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

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

    ax.set_yscale('log')
    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()