import numpy as np
from IPython import embed
from scipy.signal import convolve
import matplotlib.mlab as mlab


def avgNestedLists(nested_vals):
    """
    Averages a 2-D array and returns a 1-D array of all of the columns
    averaged together, regardless of their dimensions.
    """
    output = []
    maximum = 0
    for lst in nested_vals:
        if len(lst) > maximum:
            maximum = len(lst)
    for index in range(maximum):  # Go through each index of longest list
        temp = []
        for lst in nested_vals:  # Go through each list
            if index < len(lst):  # If not an index error
                temp.append(lst[index])
        output.append(np.nanmean(temp))
    return output


def gaussian_convolve(spike_train, fxn, sampling_rate):

    all_convolve_spikes = []
    trial_length = int((spike_train[-1] - spike_train[0]) * sampling_rate)
    spike_train = spike_train - spike_train[0]  # changing spike train to start at 0 (subtracting baseline)
    trial_time = np.arange(0, (trial_length + 1), 1)
    trial_bool = np.zeros(len(trial_time))
    #  Boolean list in length of trial length, where 1 means spike happened, 0 means no spike
    spike_indx = (spike_train * sampling_rate).astype(np.int)
    trial_bool[spike_indx] = 1
    # trial_bool = trial_bool[30000:(len(trial_bool)-30000)]
    convolve_spikes = np.asarray(
        [convolve(trial_bool, fxn, mode='valid')])  # convolve gaussian with boolean spike list
    all_convolve_spikes.append(convolve_spikes[0, :])
    avg_convolve_spikes = avgNestedLists(all_convolve_spikes)
    return avg_convolve_spikes


def fourier_psd(avg_convolve_spikes, sampling_rate):
    p, freq = mlab.psd(avg_convolve_spikes, NFFT=sampling_rate * 3, noverlap=sampling_rate * 1.5,
                       Fs=sampling_rate,
                       detrend=mlab.detrend_mean)
    std_four = np.std(freq[5:])
    mn_four = np.mean(freq)
    return p, freq, std_four, mn_four