#!/usr/bin/env python
import pandas as pd
import numpy as np
import math

###############################################################################
# Modules for calculating gev extremes
###############################################################################
 
# Calculate the extremes parameters
def get_gev(arr):
    edf = pd.DataFrame({'val': arr})
    edf.dropna(inplace=True)
    edf.sort_values(by=['val'], ascending=False, inplace=True)
    
    edf = edf.reset_index()[['val']]
    edf['Rank'] = edf.index + 1
    n = len(edf.index)

    edf['B1i'] = (n - edf['Rank']) / (n * (n - 1)) * edf['val']
    edf['B2i'] = edf['B1i'] * (n - 1 - edf['Rank']) / (n - 2)
    edf['B3i'] = edf['B2i'] * (n - 2 - edf['Rank']) / (n - 3)

    B = [0, 0, 0, 0, 0]
    L = [0, 0, 0, 0, 0]
    c = [0, 0, 0, 0]
    kappa = [0, 0, 0, 0]
    alpha = [0, 0, 0, 0]
    psi = [0, 0, 0, 0]

    gev = pd.DataFrame({"kappa": [3, 0, 0, 0],
                        "alpha": [2, 0, 0, 0],
                        "psi": [0, 0, 0, 0]})

    B[0] = edf['val'].mean()
    B[1] = edf['B1i'].sum()
    B[2] = edf['B2i'].sum()
    B[3] = edf['B3i'].sum()

    L[1] = B[0]
    L[2] = 2 * B[1] - B[0]
    L[3] = 6 * B[2] - 6 * B[1] + B[0]

    # GEV distribution using L moments
    c[0] = (2 * L[2]) / (L[3] + 3 * L[2]) - 0.630930
    gev.loc[0, 'kappa'] = 7.8590 * c[0] + 2.9554 * c[0] * c[0]
    gamma = math.gamma(1 + gev.loc[0, 'kappa'])
    gev.loc[0, 'alpha'] = gev.loc[0, 'kappa'] * L[2] / (gamma * (1 - math.pow(2, -gev.loc[0, 'kappa'])))
    gev.loc[0, 'psi'] = L[1] + gev.loc[0, 'alpha'] * (gamma - 1) / gev.loc[0, 'kappa']

    return gev


# Helper function to calculate extremes for given probability
def get_extreme(gev, prob):
    a = gev.loc[0, 'alpha']
    k = gev.loc[0, 'kappa']
    p = gev.loc[0, 'psi']

    dist = p + a / k * (1 - (np.power(-np.log(prob), k)))
    return dist


# Calculates 
def get_highs(arr, yrs=[2, 5, 10, 50, 100]):
    df = pd.DataFrame({"yrs" : yrs})
    df['prob'] = 1 - 1 / df['yrs']

    gev = get_gev(arr)
    df['dist0'] = get_extreme(gev, df['prob'])
    return df[['dist0']]
    
    #epdf['prob'] = epdf['prob'].apply(lambda x : int(round(1 / (1 - x), 0)))
    #epdf = epdf.rename(columns={'prob': 'returnYr'})


def get_lows(arr, yrs=[2, 5]):
    df = pd.DataFrame({"yrs" : yrs})
    df['prob'] = 1/ df['yrs']
    
    gev = get_gev(arr)
    df['dist0'] = get_extreme(gev, df['prob'])
    return df[['dist0']]
    
    #eldf['prob'] = eldf['prob'].apply(lambda x : int(round((1 / x), 0)))
    #eldf = eldf.rename(columns={'prob': 'returnYr'})