# AUTOGENERATED! DO NOT EDIT! File to edit: notebooks/03_add_responses.ipynb (unless otherwise specified).

__all__ = ['interpolate', 'butterFilter', 'find_first_peak', 'find_response', 'getResponses', 'add_responses']

# Cell
from .constants import Constants as C
import numpy as np # 1.18.4
import pandas as pd #'1.0.3'
import peakutils #'1.3.3'
from scipy.interpolate import splrep, splev # 1.4.1
from scipy.signal import butter, filtfilt
from tqdm import tqdm #'4.46.0'

def interpolate(times, accelerations):
    interpolated = np.array([])
    # Times sometimes comes in as a nan value requring this try except
    try:
        if len(times):
            spl = splrep(times, accelerations)
            x2 = np.arange(times[0], times[-1])
            interpolated = splev(x2, spl)
    except:
        pass#
    return interpolated


def butterFilter(acc, cutoff, frequency, order):
    nyq = 0.5 * frequency
    normal_cutoff = cutoff / nyq
    b, a = butter(order, normal_cutoff, btype='low', analog=False)
    buttered = filtfilt(b, a, acc)
    return buttered


def find_first_peak(acceleration, threshold, min_dist):
    first_peak, force, response = np.nan, np.nan, "NA" # If no acceleration present, set response to NA
    if (len(acceleration)>0):
        peaks = peakutils.peak.indexes(abs(acceleration), thres=.3, min_dist=min_dist,thres_abs=False)
        if len(peaks): # If peaks were detected
            first_peak = peaks[0]
            force = acceleration[first_peak]
            if force > threshold: # Set response to pull if bigger than threshold
                response = "pull"
            elif force < -threshold: # Set response to push if smaller than threshold
                response = "push"
            else: # If response did not pass threshold, set it to timeout
                first_peak = np.nan
                force = np.nan
                response = "none"
        else:
            response = "ND" # If no peaks are detected set response to not determined
    return (first_peak, force, response)


def find_response(response, acceleration, first_peak):
    rt = np.nan
    if not np.isnan(first_peak):
        rts = np.where((abs(acceleration[1:first_peak-10])<.5))[0]
        if len(rts):
            rt = rts[-1]
        else:
            response = "ND"
    return rt, response


def getResponses(line, min_dist, threshold, butter_filter):
    # Smooth the data to millisecond resolution
    acceleration = interpolate(line.loc[C.TIME_COLUMN], line.loc[C.ACCELERATION_COLUMN])
    # Butter filter
    if butter_filter == True:
        acceleration = butterFilter(acceleration, 10, 1000, 2)
    # Find the first peak (location and magnitude)
    first_peak, force, response = find_first_peak(acceleration, threshold, min_dist)
    # Finding the reaction time
    rt, response = find_response(response, acceleration, first_peak)
    # Reduce size of acceleration column
    acceleration = acceleration.astype('float16')
    output = [response, rt, abs(force), first_peak]
    output.append(acceleration)
    gyro = interpolate(line.loc[C.GYRO_TIME_COLUMN], line.loc[C.GYRO_Y_COLUMN])
    gyro = gyro.astype('float16')
    output.append(gyro)
    return output


def add_responses(df, min_dist = 10, threshold = 1, butter_filter = False):
    # Initiating progress bar
    #tqdm.pandas()
    statistics = [C.RESPONSE_COLUMN,
                  C.RT_COLUMN,
                  C.FORCE_COLUMN,
                  C.PEAK_AT_COLUMN,
                  C.INTERPOLATED_COLUMN,
                  C.INTERPOLATED_GYRO_COLUMN]

    df[statistics] = df.apply(getResponses,axis=1,result_type="expand", min_dist = min_dist, threshold=threshold, butter_filter=butter_filter)
    df[C.ACCURACY_COLUMN] = df.response == df.correct_response
    df[C.RT_INV_COLUMN] = 1000/df[C.RT_COLUMN]
    df = df.drop(columns=['level_0','times','accelerations','accelerations_x',
       'accelerations_y','gyro_times','gyro_x', 'gyro_y', 'gyro_z'], errors='ignore')
    # Changin null values to make operations in R easier
    categorical_question_columns = ['condition','correct_response', 'device','drawn_at','experiment','sensor_type','stimulus','stimulus_set']
    df[categorical_question_columns] = df[categorical_question_columns].where(pd.notnull(df[categorical_question_columns]),None)
    numeric_question_columns = ['drawn_at', 'drawn_at_unix']
    df[numeric_question_columns] = df[numeric_question_columns].astype(float)
    df['is_pull'] = df.response == 'pull'
    return df.reset_index(drop = False)