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(all_spike_trains, fxn, sampling_rate, time):
    """
    Takes an array of spike trains of different sizes,
    convolves it with a gaussian, returns the average gaussian convolve spikes
    """
    all_convolve_spikes = []
    for spike_train in all_spike_trains:
        time_cutoff = time * sampling_rate
        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, :][:time_cutoff])
    #
    # cutoff = min([len(i) for i in all_convolve_spikes])
    # for ix, convolved in enumerate(all_convolve_spikes):
    #     all_convolve_spikes[ix] = all_convolve_spikes[ix][:cutoff]
    #avg_convolve_spikes = avgNestedLists(all_convolve_spikes)
    avg_convolve_spikes = np.mean(all_convolve_spikes, 0)
    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