'''
Spatio-temporal animation
'''
from brian2 import ms
from matplotlib.animation import FuncAnimation
import numpy as np
import matplotlib.pyplot as plt



def create_neural_sheet_animation(times, activities, frame_rate=25):
    """

    :param times: array of time points
    :param activtivies: 3d array of spiking activity from zero no spike to one very recent spike with shape(#neuros,
    #neurons,
    time)
    :param frame_rate: in hertz
    :return: animation
    """
    number_of_frames = times.shape[0]

    fig, ax = plt.subplots(1,1)
    image = ax.imshow(activities[:,:,0], vmin=0, vmax=1, cmap="gray")
    timer = ax.set_title("t={:10.1f}ms".format(times[0]))

    def animate(frame):
        image.set_array(activities[:,:,frame])
        timer.set_text("t={:10.1f}ms".format(times[frame]))
        return image

    animation = FuncAnimation(fig, animate, number_of_frames, interval=1000/frame_rate)
    return animation

#number_of_excitatory_neurons_per_row = 30
#times = np.arange(0, 10, 0.5)

#activities = np.random.choice(np.array([0,1]),(number_of_excitatory_neurons_per_row,
 #                                             number_of_excitatory_neurons_per_row,
  #                         times.shape[0]))

#animation = create_neural_sheet_animation(times, activities)

#plt.show()
#animation.save("out_2dgrid.mp4", dpi=number_of_excitatory_neurons_per_row, writer='imagemagick')

def transform_spike_train_in_activity_array(spike_trains, dt_binning, length, activity_decay_time_constant):
    times = np.arange(0, length / ms, dt_binning / ms) * ms
    number_of_neurons = len(spike_trains)
    number_of_neurons_per_row = int(np.sqrt(number_of_neurons))
    spike_activities = [transform_spike_times_to_activity(spike_times, dt_binning, length,
                                                          activity_decay_time_constant)
                        for spike_times in spike_trains]

    spike_activities = np.reshape(np.array(spike_activities), (number_of_neurons_per_row,
                                                           number_of_neurons_per_row,-1))
    return times, spike_activities


def transform_spike_times_to_activity(spike_times, dt_binning, length, activity_decay_time_constant):
    kernel_times = np.arange(0, 4 * activity_decay_time_constant / ms, dt_binning / ms) * ms
    activity_kernel = np.exp(-kernel_times / activity_decay_time_constant)
    number_of_bins = int(length/dt_binning)+1
    indices_with_a_spike = np.floor(spike_times / dt_binning).astype(int)
    spike_train_in_binary_form = np.zeros((number_of_bins,))
    spike_train_in_binary_form[indices_with_a_spike] = 1
    spike_activity = np.convolve(spike_train_in_binary_form, activity_kernel, mode='same')
    return spike_activity