# -*- coding: utf-8 -*-

"""
Author: Rien Sonck
Date: 2020-06-26
Description: This file contains all the functions that I use to analyse the simulated data from the sync model. 
"""

######################
# Importing packages #
######################
import numpy as np
import pandas as pd
from scipy import ndimage
from Sync_Model import model_sim
from matplotlib import pyplot as pl

###############
#  Functions  #
###############
def read_in_data(mfc_freq, thres, path):
    """ 
    Reading in data from files generated by the sync model. 

    Parameters
    ----------
    mfc_freq : numpy array 
               array of mfc frequency values
    thres    : numpy array
               array of threshold values
    path     : string
               path to the folder containing the simulated data
    """ 
    data = []
    # reading in the data
    for mfc in mfc_freq: 
        for thr in thres:
            file = '{0}Behavioral_Data_simulation_sub0_thetaFreq{1}.00Hz_thresh{2}_drift2.0.csv'.format(path, mfc, thr)
            subset_data = pd.read_csv(file, index_col=None, header=0)
            subset_data['mfc'] = mfc
            subset_data['thres'] = thr
            data.append(subset_data)
    # creating the pandas dataframe
    pd_data = pd.concat(data, axis=0, ignore_index=True)
    return(pd_data)

def plot_double_RT(pd_subset, ax):
    """ 
    Plotting a single histogram plot. 

    Parameters
    ----------
    pd_subset : pandas dataframe 
                a subset of data from pd_data, this subset has been selected based on one mfc and one threshold value.
    ax        : matplotlib ax
                the ax on which to plot the figure 
    """ 

    # making the error RTs negative
    pd_subset.loc[pd_subset['accuracy'] == 0, 'rt'] = pd_subset.loc[pd_subset['accuracy'] == 0]['rt'].values * -1
    # extracting RT
    rt_err = pd_subset.loc[pd_subset['accuracy'] == 0, 'rt'].values
    rt_cor = pd_subset.loc[pd_subset['accuracy'] == 1, 'rt'].values
    # Creating bins
    bins_err = np.arange(-700, 0.1, 25) 
    bins_cor = np.arange(0, 700.1, 25)
    # getting histogram data
    heights_err, bins_err = np.histogram(rt_err, bins=bins_err)
    heights_cor, bins_cor = np.histogram(rt_cor, bins=bins_cor)
    # calculating the overall percantage 
    percent_err = [i / (sum(heights_err) + sum(heights_cor)) * 100 for i in heights_err]
    percent_cor = [i / (sum(heights_err) + sum(heights_cor)) * 100 for i in heights_cor]
    # plotting the histogram
    ax.bar(bins_cor[:-1]+12.5, percent_cor, width=25, edgecolor='black', label='Correct')
    ax.bar(bins_err[1:]-12.5, percent_err, width=25, edgecolor='black', label='Error')
    # styling
    vals = np.arange(0, 16 ,5)
    ax.set_yticks(vals) 
    ax.set_yticklabels(['%1.f%%' %i for i in vals])
    ax.set_xlabel("reaction time"); ax.set_ylabel("percentage")
    ax.set_title("Threshold: {0}, MFC Freq: {1}".format(pd_subset["thres"].unique()[0], pd_subset["mfc"].unique()[0]))
    ax.legend(loc='best')

def plot_RT_difference(pd_subset, pd_targetData, ax):
    """ 
    Plotting a single histogram plot. 

    Parameters
    ----------
    pd_subset     : pandas dataframe 
                    a subset of data from pd_data, this subset has been selected based on one mfc and one threshold value
    pd_targetData : pandas dataframe
                    dataframe containing the target data
    ax            : matplotlib ax
                    the ax on which to plot the figure 
    """ 
    # extracting RT
    rt_err = pd_subset.loc[pd_subset['accuracy'] == 0, 'rt'].values * -1
    rt_cor = pd_subset.loc[pd_subset['accuracy'] == 1, 'rt'].values
    rt_err_target = pd_targetData.loc[pd_targetData['accuracy'] == 0, 'rt'].values * -1
    rt_cor_target = pd_targetData.loc[pd_targetData['accuracy'] == 1, 'rt'].values
    # creating bins
    bins_err = np.arange(-700, 0.1, 25) 
    bins_cor = np.arange(0, 700.1, 25)
    # getting histogram data
    heights_err, bins_err = np.histogram(rt_err, bins=bins_err)
    heights_cor, bins_cor = np.histogram(rt_cor, bins=bins_cor)
    heights_err_target, bins_err = np.histogram(rt_err_target, bins=bins_err)
    heights_cor_target, bins_cor = np.histogram(rt_cor_target, bins=bins_cor)
    # calculating the overall percantage 
    percent_err = np.array([i / (sum(heights_err) + sum(heights_cor)) * 100 for i in heights_err])
    percent_cor = np.array([i / (sum(heights_err) + sum(heights_cor)) * 100 for i in heights_cor])
    percent_err_target = np.array([i / (sum(heights_err_target) + sum(heights_cor_target)) * 100 for i in heights_err_target])
    percent_corr_target = np.array([i / (sum(heights_err_target) + sum(heights_cor_target)) * 100 for i in heights_cor_target])
    # calculating the difference
    diff_err = percent_err_target - percent_err
    diff_cor = percent_corr_target - percent_cor
    squared_err = np.round(sum(diff_err**2) + sum(diff_cor**2), 2) 
    # plotting the histogram
    ax.bar(bins_cor[:-1]+12.5, diff_cor, width=25, edgecolor='black', label='Correct')
    ax.bar(bins_err[1:]-12.5, diff_err , width=25, edgecolor='black', label='Error')
    # styling
    vals = np.arange(-10, 11 , 10)
    ax.set_yticks(vals) 
    ax.set_xlabel("reaction time"); ax.set_ylabel("difference error")
    ax.set_title("Threshold: {0}, MFC Freq: {1}\nSquared Error: {2}".format(
                                        pd_subset["thres"].unique()[0], 
                                        pd_subset["mfc"].unique()[0],
                                        squared_err))
    ax.legend(loc='best')

def goal_func(x, *args):
    ######################################
    # Value Extraction and Initalization #
    ######################################
    """ 
    function to be minimized by the differential evolution algorithm 

    Parameters
    ----------
    x    : integer
           this value lies in-between the specified bounds, and will be generated by the differential evolution algorithm. 
    args : tuple
           any additional fixed model parameters and other variables needed to completely specify the objective function.
    """ 
    gd = args[0]              # gamma distance model parameter
    thres = args[1]           # mfc frequency model parameter
    pd_targetData = args[2]   # target dataframe
    squared_err = []
    ################################
    # Generating model simulations #
    ################################
    pd_data = model_sim(Threshold = thres, drift = gd, theta_freq = x, save_behavioral = False, return_dataframe = True)
    #################################
    # Calculating the Squared Error #
    #################################
    accuracy = pd_data['accuracy'].unique() 
    for acc in accuracy: 
        # extracting RTs
        rt_sim = pd_data.loc[pd_data['accuracy'] == acc, 'rt'].values               # simulated RTs
        rt_target = pd_targetData.loc[pd_targetData['accuracy'] == acc, 'rt'].values    # target RTs
        # creating bins
        bins = np.arange(0, 700.1, 25)
        # getting histogram data
        heights_sim, bins = np.histogram(rt_sim, bins=bins)
        heights_target, bins = np.histogram(rt_target, bins=bins)
        # calculating the difference
        diff = heights_target - heights_sim
        # calculating the squared error
        err = sum(diff**2)
        squared_err.append(err)
    # squared_err[1] is error, squared_err[0] is correct
    error = squared_err[1] + squared_err[0]
    print("Trying MFC: {0}, error: {1}".format(x, error))
    return error

def new_goal_func(x, *args):
    ######################################
    # Value Extraction and Initalization #
    ######################################
    """ 
    function that returns the error of the goal function based on previously generated 
    simulated data files by Sync_Model_Simulations.py

    Parameters
    ----------
    x    : integer
           this value lies in-between the specified bounds, and will be generated by the differential evolution algorithm. 
    args : tuple
           any additional fixed model parameters and other variables needed to completely specify the objective function.
    """ 
    gd = args[0]              # gamma distance model parameter
    thres = args[1]           # mfc frequency model parameter
    target_file = args[2]     # target file
    squared_err = []
    ####################
    # Reading in data #
    ###################
    # reading in the target data
    pd_targetData = pd.read_csv(target_file)
    ################################
    # Generating model simulations #
    ################################
    pd_data = model_sim(Threshold = thres, drift = gd, theta_freq = x, save_behavioral = False, return_dataframe = True)
    ###################################
    # Calculating RT Shape difference #
    ###################################
    accuracy = pd_data['accuracy'].unique() 
    for acc in accuracy: 
        # extracting RTs
        rt_sim = pd_data.loc[pd_data['accuracy'] == acc, 'rt'].values               # simulated RTs
        rt_target = pd_targetData.loc[pd_targetData['accuracy'] == acc, 'rt'].values    # target RTs
        # creating bins
        bins = np.arange(0, 700.1, 25)
        # getting histogram data
        heights_sim, bins = np.histogram(rt_sim, bins=bins)
        heights_target, bins = np.histogram(rt_target, bins=bins)
        # calculating the difference
        diff = heights_target - heights_sim
        # calculating the squared error
        err = sum(diff**2)
        squared_err.append(err)

    ###################################
    # Calculating RT Peak difference  #
    ###################################
    # applying kernels to the error and correct RT distributions
    # getting the right sigma for smoothing
    cor_indices_target, err_indices_target = [], []
    cor_indices_sim, err_indices_sim = [], []
    bins = np.arange(0, 1000, 10)

    sigma_cor = 5
    sigma_err = 5
    # this while loop will keep reducing the sigma until we can recover the second peak
    while len(cor_indices_target) < 2 or len(err_indices_target) < 2:
        if sigma_err == 0 or sigma_cor == 0:
            error = 50000
            return error
        RT_hist_list_target = RT_distribution_smoothing(pd_targetData, bins, sigma_err, sigma_cor)
        # finding where in the bins the peaks and troughs happen and mark it 
        time_step_change_target = np.diff(np.sign(np.diff(RT_hist_list_target)))
        # extracting the index of the peak and trough bins based on the marker
        cor_indices_target = np.where((time_step_change_target[1]!= 0) & (time_step_change_target[1] != 1))[0]
        err_indices_target = np.where((time_step_change_target[0] != 0) & (time_step_change_target[0] != 1))[0] 
        if len(cor_indices_target) < 2: 
            sigma_cor = sigma_cor - 1
        if len(err_indices_target) < 2:
            sigma_err = sigma_err - 1

    sigma_cor = 5
    sigma_err = 5
    while len(cor_indices_sim) < 2 or len(err_indices_sim) < 2:
        if sigma_err == 0 or sigma_cor == 0:
            error = 50000
            return error
        RT_hist_list_sim = RT_distribution_smoothing(pd_data, bins, sigma_err, sigma_cor)
        # finding where in the bins the peaks and troughs happen and mark it 
        time_step_change_sim = np.diff(np.sign(np.diff(RT_hist_list_sim)))
        # extracting the index of the peak and trough bins based on the marker
        cor_indices_sim = np.where((time_step_change_sim[1]!= 0) & (time_step_change_sim[1] != 1))[0]
        err_indices_sim = np.where((time_step_change_sim[0] != 0) & (time_step_change_sim[0] != 1))[0]
        if len(cor_indices_sim) < 2: 
            sigma_cor = sigma_cor - 1
        if len(err_indices_sim) < 2:
            sigma_err = sigma_err - 1

    # extracting the peaks and trough bin value
    peaks_troughs_err_target = RT_hist_list_target[0][err_indices_target]
    peaks_troughs_cor_target = RT_hist_list_target[1][cor_indices_target]
    peaks_troughs_err_sim = RT_hist_list_sim[0][err_indices_sim]
    peaks_troughs_cor_sim = RT_hist_list_sim[1][cor_indices_sim]
    # calculate the relative peak height
    relative_height_cor_target = peaks_troughs_cor_target[1] / peaks_troughs_cor_target[0]
    relative_height_err_target = peaks_troughs_err_target[1] / peaks_troughs_err_target[0]
    relative_height_cor_sim = peaks_troughs_cor_sim[1] / peaks_troughs_cor_sim[0] 
    relative_height_err_sim = peaks_troughs_err_sim[1] / peaks_troughs_err_sim[0] 
    # error
    peak_error = np.abs(relative_height_cor_sim - relative_height_cor_target) + np.abs(relative_height_err_sim - relative_height_err_target)

    # squared_err[1] is error, squared_err[0] is correct
    # peak_error * 10**6 -
    error = squared_err[1] + squared_err[0] + (peak_error * 10**5)

    #print("mfc: {0}, thres: {1}, gd: {2}, error: {3}".format(x, thres, gd, error))
    return error

def goal_func_from_files(mfc, thres, gd, pd_targetData, path):
    ######################################
    # Value Extraction and Initalization #
    ######################################
    """ 
    function that returns the error of the goal function based on previously generated 
    simulated data files by Sync_Model_Simulations.py

    Parameters
    ----------
    mfc           : integer
                    the mfc of the simulated data file
    thres         : integer
                    the thres of the simulated data file
    gd            : integer
                    the gamma distance of the simulated data file
    pd_targetData : pandas dataframe
                    containing the data of the target data file
    path          : string
                    containing the path to the folder where all the data files are stored
    """ 
    squared_err = []
    ####################
    # Reading in data #
    ###################
    target = '{0}Behavioral_Data_simulation_sub0_thetaFreq{1}.00Hz_thresh{2}_drift{3}.0.csv'.format(path, mfc, thres, gd)
    pd_data = pd.read_csv(target)
    #################################
    # Calculating the Squared Error #
    #################################
    accuracy = pd_data['accuracy'].unique() 
    for acc in accuracy: 
        # extracting RTs
        rt_sim = pd_data.loc[pd_data['accuracy'] == acc, 'rt'].values               # simulated RTs
        rt_target = pd_targetData.loc[pd_targetData['accuracy'] == acc, 'rt'].values    # target RTs
        # creating bins
        bins = np.arange(0, 700.1, 25)
        # getting histogram data
        heights_sim, bins = np.histogram(rt_sim, bins=bins)
        heights_target, bins = np.histogram(rt_target, bins=bins)
        # calculating the difference
        diff = heights_target - heights_sim
        # calculating the squared error
        err = sum(diff**2)
        squared_err.append(err)
    # squared_err[1] is error, squared_err[0] is correct
    error = squared_err[1] + squared_err[0]
    print("Trying MFC: {0}, thres: {1}, gd: {2} error: {3}".format(mfc, thres, gd, error))
    return error

def new_goal_func_from_files(mfc, thres, gd, target_file, path):
    ######################################
    # Value Extraction and Initalization #
    ######################################
    """ 
    function that returns the error of the goal function based on previously generated 
    simulated data files by Sync_Model_Simulations.py

    Parameters
    ----------
    mfc           : integer
                    the mfc of the simulated data file
    thres         : integer
                    the thres of the simulated data file
    gd            : integer
                    the gamma distance of the simulated data file
    target_path   : string
                    containing the path to the the target file
    path          : string
                    containing the path to the folder where all the data files are stored
    """ 
    squared_err = []
    ####################
    # Reading in data #
    ###################
    # reading in the target data
    pd_targetData = pd.read_csv(target_file)
    # reading in the simulated data
    simulation_file = '{0}Behavioral_Data_simulation_sub0_thetaFreq{1:.2f}Hz_thresh{2}_drift{3}.0.csv'.format(path, mfc, thres, gd)
    pd_data = pd.read_csv(simulation_file)
    ###################################
    # Calculating RT Shape difference #
    ###################################
    accuracy = pd_data['accuracy'].unique() 
    for acc in accuracy: 
        # extracting RTs
        rt_sim = pd_data.loc[pd_data['accuracy'] == acc, 'rt'].values               # simulated RTs
        rt_target = pd_targetData.loc[pd_targetData['accuracy'] == acc, 'rt'].values    # target RTs
        # creating bins
        bins = np.arange(0, 700.1, 25)
        # getting histogram data
        heights_sim, bins = np.histogram(rt_sim, bins=bins)
        heights_target, bins = np.histogram(rt_target, bins=bins)
        # calculating the difference
        diff = heights_target - heights_sim
        # calculating the squared error
        err = sum(diff**2)
        squared_err.append(err)

    ###################################
    # Calculating RT Peak difference  #
    ###################################
    # applying kernels to the error and correct RT distributions
    # getting the right sigma for smoothing
    cor_indices_target, err_indices_target = [], []
    cor_indices_sim, err_indices_sim = [], []
    bins = np.arange(0, 1000, 10)

    sigma_cor = 5
    sigma_err = 5

    RT_hist_list_target = RT_distribution_smoothing(pd_targetData, bins, sigma_err, sigma_cor)
    RT_hist_list_sim = RT_distribution_smoothing(pd_data, bins, sigma_err, sigma_cor)
    time_step_change_target = np.diff(np.sign(np.diff(RT_hist_list_target)))
    time_step_change_sim = np.diff(np.sign(np.diff(RT_hist_list_sim)))
    cor_indices_target = np.where((time_step_change_target[1]!= 0) & (time_step_change_target[1] != 1))[0]
    err_indices_target = np.where((time_step_change_target[0] != 0) & (time_step_change_target[0] != 1))[0] 
    cor_indices_sim = np.where((time_step_change_sim[1]!= 0) & (time_step_change_sim[1] != 1))[0]
    err_indices_sim = np.where((time_step_change_sim[0] != 0) & (time_step_change_sim[0] != 1))[0]

    """
    # this while loop will keep reducing the sigma until we can recover the second peak
    while len(cor_indices_target) < 2 or len(err_indices_target) < 2:
        RT_hist_list_target = RT_distribution_smoothing(pd_targetData, bins, sigma_err, sigma_cor)
        # finding where in the bins the peaks and troughs happen and mark it 
        time_step_change_target = np.diff(np.sign(np.diff(RT_hist_list_target)))
        # extracting the index of the peak and trough bins based on the marker
        cor_indices_target = np.where((time_step_change_target[1]!= 0) & (time_step_change_target[1] != 1))[0]
        err_indices_target = np.where((time_step_change_target[0] != 0) & (time_step_change_target[0] != 1))[0] 
        if len(cor_indices_target) < 2: 
            sigma_cor = sigma_cor - 1
        if len(err_indices_target) < 2:
            sigma_err = sigma_err - 1


    sigma_cor = 5
    sigma_err = 5
    while len(cor_indices_sim) < 2 or len(err_indices_sim) < 2:
        RT_hist_list_sim = RT_distribution_smoothing(pd_data, bins, sigma_err, sigma_cor)
        # finding where in the bins the peaks and troughs happen and mark it 
        time_step_change_sim = np.diff(np.sign(np.diff(RT_hist_list_sim)))
        # extracting the index of the peak and trough bins based on the marker
        cor_indices_sim = np.where((time_step_change_sim[1]!= 0) & (time_step_change_sim[1] != 1))[0]
        err_indices_sim = np.where((time_step_change_sim[0] != 0) & (time_step_change_sim[0] != 1))[0]
        if len(cor_indices_sim) < 2: 
            sigma_cor = sigma_cor - 1
        if len(err_indices_sim) < 2:
            sigma_err = sigma_err - 1
    """

    # extracting the peaks and trough bin value
    peaks_troughs_err_target = RT_hist_list_target[0][err_indices_target]
    peaks_troughs_cor_target = RT_hist_list_target[1][cor_indices_target]
    peaks_troughs_err_sim = RT_hist_list_sim[0][err_indices_sim]
    peaks_troughs_cor_sim = RT_hist_list_sim[1][cor_indices_sim]
    # calculate the relative peak height
    #relative_height_cor_target = peaks_troughs_cor_target[1] / peaks_troughs_cor_target[0]
    #relative_height_err_target = peaks_troughs_err_target[1] / peaks_troughs_err_target[0]
    #relative_height_cor_sim = peaks_troughs_cor_sim[1] / peaks_troughs_cor_sim[0] 
    #relative_height_err_sim = peaks_troughs_err_sim[1] / peaks_troughs_err_sim[0]
    relative_height_target = peaks_troughs_err_target[0] #- peaks_troughs_cor_target[0]
    relative_height_sim = peaks_troughs_err_sim[0] #- peaks_troughs_cor_sim[0]
    # error
    #peak_error = np.abs(relative_height_cor_sim - relative_height_cor_target) + np.abs(relative_height_err_sim - relative_height_err_target)
    peak_error =  abs(relative_height_sim / relative_height_target)

    # squared_err[1] is error, squared_err[0] is correct
    # peak_error * 10**6 -
    #error = squared_err[1] + squared_err[0] + (peak_error * 10**5)
    error = (peak_error * 10**2)

    print("MFC: {0}, thres: {1}, gd: {2}, peak_error_n: {3}, peak_error:{4} RT error: {5}, RT correct: {6},error: {7}".format(mfc, thres, gd, peak_error, peak_error * 10**4, squared_err[1], squared_err[0], error))
    return error

def RT_distribution_smoothing(pd_data, bins, sigma_err, sigma_cor):
    """ 
    Function that applies the difference kernel and then a Gaussian smoothing kernel to the error 
    and correct RT distributions corresponding to each MFC frequency in the model.  

    Parameters
    ----------
    pd_data    : pandas dataframe 
                 dataframe containing the data
    bins       : numpy array
                 array containing the bins for the histograms
    sigma      : integer
                 smoothing of the Gaussian kernel
    """ 
    accuracy = [0, 1]
    thetas_hist_list = []
    theta_hist = np.zeros(shape=[2, 1, len(bins)-1])
    ###################
    # Reading in data #
    ###################
    # creating histograms
    theta_hist[0, :] = np.histogram(pd_data.loc[pd_data['accuracy']== 0, 'rt'].values, bins=bins, normed = True)[0] # error histograms
    theta_hist[1, :] = np.histogram(pd_data.loc[pd_data['accuracy']== 1, 'rt'].values, bins=bins, normed = True)[0] # correct histograms
    ###################################
    # Difference Kernel and Smoothing #
    ###################################
    # run difference kernel
    theta_histDiff_err = -np.array([theta_hist[0][0][bin_n] - (theta_hist[0][0][bin_n-1] + theta_hist[0][0][bin_n+1]).mean() for bin_n in np.arange(1, len(bins)-2)]).T
    theta_histDiff_cor = -np.array([theta_hist[1][0][bin_n] - (theta_hist[1][0][bin_n-1] + theta_hist[1][0][bin_n+1]).mean() for bin_n in np.arange(1, len(bins)-2)]).T
    # smooth
    theta_histDiffSmooth_err = ndimage.gaussian_filter1d(theta_histDiff_err, axis = 0, sigma = sigma_err)
    theta_histDiffSmooth_cor = ndimage.gaussian_filter1d(theta_histDiff_cor, axis = 0, sigma = sigma_cor)
    pl.plot(theta_histDiffSmooth_err)
    pl.plot(theta_histDiffSmooth_cor)
    return (theta_histDiffSmooth_err, theta_histDiffSmooth_cor)

################
#  Main Code   #
################
"""
 __name__ of this file will be assigned 'Analyse_Functions' when importing it to another file.
 The code below will make sure that the other file doesn't execute the functions in this file while
 importing this file (http://effbot.org/pyfaq/tutor-what-is-if-name-main-for.htm)
"""
if __name__ == 'Analyse_Functions':
    pass