# -*- coding: utf-8 -*-
"""
Created on Wed Mar  4 09:50:31 2020

@author: Johan Rene van Dorp
"""

import numpy as np
import matplotlib.patches as patches
import pandas as pd
from scipy.stats import linregress

def Estimate_empirical_pmf(outcomes,samples):
    # Counting the number of times a sample appears
    counts = dict()
    for sample in samples:
        if sample not in counts:
            counts[sample] = 1
        else:
            counts[sample] = counts[sample]+1
    # Adding the outcomes to the counts that were not present in the sample        
    for outcome in outcomes:
        if outcome not in counts:
            counts[outcome] = 0
    # evaluating the counts in the order of the outcomes        
    lst = list(counts.keys()) 
    lst.sort()
    the_counts = np.zeros(len(outcomes))
    i = -1
    for key in lst:
        i = i + 1
        the_counts[i] = counts[key]
    # Finally converting the counts to emprical pmf values    
    pmf = the_counts / sum(the_counts)  
    return(pmf)

def Calc_boundaries_and_widths(outcomes): 
    # This functions determines the bounds and width for the plot of
    # a pmf with the outcomes as the possible occurrences
    mid_points = (outcomes[0:len(outcomes)-1]+outcomes[1:len(outcomes)])/2
    first_element = outcomes[0]
    last_element = outcomes[len(outcomes)-1]
    bounds = np.append(first_element,mid_points)
    bounds = np.append(bounds,last_element) 
    widths = bounds[1:len(bounds)]-bounds[0:len(bounds)-1]
    return(bounds,widths)

def add_pmf_prob_hist_to_figure_panel(figure_panel,outcomes,probs,color,the_alpha):
    bounds, widths = Calc_boundaries_and_widths(outcomes)
    for i in np.arange(len(bounds)-1):
        figure_panel.add_patch(patches.Rectangle(
                          (bounds[i], 0), 
                           widths[i], 
                           probs[i],
                           facecolor = color,
                           edgecolor = 'black',
                           alpha = the_alpha))

def add_pmf_density_hist_to_figure_panel(figure_panel,bounds,density,color):
    n = len(bounds)
    widths = bounds[1:n]-bounds[0:(n-1)]
    for i in np.arange(len(bounds)-1):
        figure_panel.add_patch(patches.Rectangle(
                          (bounds[i], 0), 
                           widths[i], 
                           density[i+1],
                           facecolor = color,
                           edgecolor = 'black',
                           alpha = 0.5))
        
def Estimate_empirical_histogram(bounds,samples): 
    n = len(bounds)
    m = len(samples)
    empirical_pmf = np.zeros(n)
    df = pd.DataFrame({'The_Sample':samples})
    cdf_lb = df[df<bounds[0]].count()/m
    for i in np.arange(1,n-1):
        cdf_ub = df[df<bounds[i]].count()/m
        width = bounds[i]-bounds[i-1]
        empirical_pmf[i] = (cdf_ub-cdf_lb)/width
        cdf_lb = cdf_ub
    return empirical_pmf        
        
def Estimate_empirical_histogram_table(bounds,samples): 
    n = len(bounds)
    m = len(samples)
    df = pd.DataFrame({'The_Sample':samples})
    cum_counts = np.zeros(n)
    counts = np.zeros(n)
    bin_prob = np.zeros(n)
    bin_pdf = np.zeros(n)
    bin_cdf = np.zeros(n)
    cum_counts[0] = df[df<=bounds[0]].count()
    for i in np.arange(1,n):
        cum_counts[i] = df[df<=bounds[i]].count()
        bin_cdf[i] = cum_counts[i]/m
        counts[i] = cum_counts[i]-cum_counts[i-1]
        bin_prob[i] = counts[i]/m
        bin_pdf[i] = bin_prob[i]/(bounds[i]-bounds[i-1])
        LB = bounds[0:n-1]
        UB = bounds[1:n]
        output_df = pd.DataFrame({'LB':LB,'UB':UB,
                          'Counts':counts[1:n],
                          'Bin Prob':bin_prob[1:n],
                          'Bin PDF':bin_pdf[1:n],
                          'Cum. Counts':cum_counts[1:n],
                          'Bin CDF':bin_cdf[1:n]})  
    return output_df

def plot_discrete_cdf(figure_panel,x,p,the_col,off_set):
    figure_panel.axhline(1, lw=1, ls = ':',color = 'darkgray',alpha = 0.5)
    # calculating cumulative probabilities
    cum_p = p
    for i in range(1,len(p)):
        cum_p[i] = cum_p[i-1]+p[i] 
    x = np.append(min(x)-off_set,x)
    x = np.append(x,max(x)+off_set)    
    y = np.append(0,cum_p)
    y = np.append(y,1)
    # plotting closed points
    closed_x = x
    closed_y = y
    open_x = x[1:len(x)-1]
    open_y = y[0:len(y)-2]
    figure_panel.scatter(closed_x,closed_y,color = the_col)
    figure_panel.scatter(open_x,open_y,edgecolor = the_col,facecolor = 'white') 
    figure_panel.vlines(open_x, open_y, closed_y[1:len(closed_y)-1], color='black', linestyles=':', lw=1)
    figure_panel.hlines(closed_y[0:len(closed_y)-1], closed_x[0:len(closed_x)-1],
                        closed_x[1:len(closed_x)], linestyles='-', lw=2,color = the_col)
    
def rsquared(x, y):
    """ Return R^2 where x and y are array-like."""

    slope, intercept, r_value, p_value, std_err = linregress(x, y)
    return r_value**2

def empirical_quantile(x,q_level):
    n_data = len(x)
    x = np.sort(x)
    F = np.arange(1,(n_data+1))/n_data
    q_index = np.sum(F<q_level)
    alpha = n_data*q_level-q_index
    x_k = x[q_index-1]
    x_k_plus_1 = x[q_index]
    emp_quantile = alpha*(x_k_plus_1-x_k)+x_k
    return emp_quantile, q_index, x_k, x_k_plus_1

def Running_Average(samples,x): 
    n = len(samples)
    r_average = np.zeros(n)
    r_stdev = np.zeros(n)
    r_cdf = np.zeros(n)
    for i in np.arange(1,n):
        r_average[i] = np.mean(samples[:i])
        r_stdev[i] = np.std(samples[:i])
        r_cdf[i] = np.count_nonzero(samples[:i] <= x)/i
    output_df = pd.DataFrame({'r_ave':r_average,
                              'r_sd':r_stdev,
                              'r_cdf':r_cdf}) 
    return output_df