import numpy as np
from brian2.units import Hz
from pypet import Trajectory
from pypet.brian2 import Brian2MonitorResult, Brian2Result


from scripts.spatial_network.head_direction_index_over_noise_scale import calculate_rates
from scripts.spatial_network.run_entropy_maximisation_orientation_map import DATA_FOLDER, TRAJ_NAME

n_exc_test_print = 465

def get_spike_train_dictionary(number_of_neurons, spike_times, neuron_indices):
    spike_train_dict = {}
    for neuron_idx in range(number_of_neurons):
        spike_train_dict[neuron_idx] = []

    for neuron_idx, t in zip(neuron_indices, spike_times):
        spike_train_dict[neuron_idx].append(t)
    return spike_train_dict


def get_firing_rate_dict_per_cell_and_direction(traj, run_name):
    traj.f_set_crun(run_name)
    firing_rate_dict = {}
    direction_names = ["dir{:d}".format(idx) for idx in range(traj.input.number_of_directions)]
    for idx in range(traj.N_E):
        firing_rate_dict[idx] = []
    for direction in direction_names:
        number_of_neurons = traj.N_E
        all_spike_times = traj.results.runs[run_name][direction].spikes.e.t
        neuron_indices = traj.results.runs[run_name][direction].spikes.e.i
        ex_spike_trains = get_spike_train_dictionary(number_of_neurons, all_spike_times,
                                                     neuron_indices)
        # print(ex_spike_trains)
        # print(ex_spike_trains[n_exc_test_print])
        ex_spike_rates = calculate_rates(ex_spike_trains.values())
        for idx, spike_rate in enumerate(ex_spike_rates):
            firing_rate_dict[idx].append(spike_rate)
    traj.f_restore_default()
    return firing_rate_dict

def get_firing_rate_array_per_cell_and_direction(traj, run_name):
    traj.f_set_crun(run_name)
    firing_rate_array = np.ndarray((traj.N_E, traj.input.number_of_directions))
    direction_names = ["dir{:d}".format(idx) for idx in range(traj.input.number_of_directions)]
    for dir_idx, direction in enumerate(direction_names):
        number_of_neurons = traj.N_E
        all_spike_times = traj.results.runs[run_name][direction].spikes.e.t
        neuron_indices = traj.results.runs[run_name][direction].spikes.e.i
        ex_spike_trains = get_spike_train_dictionary(number_of_neurons, all_spike_times,
                                                     neuron_indices)
        ex_spike_rates = calculate_rates(ex_spike_trains.values())
        for n_idx, spike_rate in enumerate(ex_spike_rates):
            #TODO: Why on earth does the unit vanish?
            firing_rate_array[n_idx, dir_idx] = spike_rate
    traj.f_restore_default()
    return firing_rate_array


def get_head_direction_indices(directions, firing_rate_array):
    n_exc_neurons = firing_rate_array.shape[0]
    n_directions = len(directions)

    tuning_vectors = np.zeros((n_exc_neurons,n_directions,2))
    # print('before:\n', tuning_vectors[9,:,:])
    for ex_id, ex_rates in enumerate(firing_rate_array):
        rate_sum = 0.
        for dir_id, dir in enumerate(directions):
            tuning_vectors[ex_id, dir_id] = np.array([np.cos(dir), np.sin(dir)]) * ex_rates[dir_id]
            rate_sum += ex_rates[dir_id]
            # if ex_id == 9:
            #     print(np.array([np.cos(dir), np.sin(dir)]) * ex_rates[dir_id])
        if rate_sum != 0.:
            tuning_vectors[ex_id, :, :] /= rate_sum
    # print('after:\n', tuning_vectors[9, :, :])
    tuning_vectors_summed = np.sum(tuning_vectors, axis=1)
    # print(tuning_vectors_summed.shape)

    head_direction_indices = np.array([np.linalg.norm(v) for v in tuning_vectors_summed])


    # print(head_direction_indices)
    return head_direction_indices


def get_runs_with_circular_morphology(traj):
    filtered_indices = traj.f_find_idx(('parameters.long_axis', 'parameters.short_axis'), lambda r1, r2: r1 == r2)
    return filtered_indices


if __name__ == "__main__":
    traj = Trajectory(TRAJ_NAME, add_time=False, dynamic_imports=Brian2MonitorResult)
    NO_LOADING = 0
    FULL_LOAD = 2
    # TODO: Why no loading of results? Because of dynamic loading?
    traj.f_load(filename=DATA_FOLDER + TRAJ_NAME + ".hdf5", load_parameters=FULL_LOAD, load_results=NO_LOADING)

    traj.v_auto_load = True

    correlation_lengths = traj.f_get('correlation_length').f_get_range()
    long_axis = traj.f_get('long_axis').f_get_range()
    short_axis = traj.f_get('short_axis').f_get_range()
    # TODO: Again, maybe directions as parameters
    directions = np.linspace(-np.pi, np.pi, traj.input.number_of_directions)

    circular_indices = list(get_runs_with_circular_morphology(traj))

    for idx, run_name in enumerate(traj.f_get_run_names()):
        firing_rate_array = get_firing_rate_array_per_cell_and_direction(traj, run_name)
        firing_rate_dict = get_firing_rate_dict_per_cell_and_direction(traj, run_name)
        traj.f_set_crun(run_name)
        print(firing_rate_array)
        print(firing_rate_dict)
        traj.f_add_result('runs.$.firing_rate_array', firing_rate_array,
                          comment='The firing rates of the excitatory population')
        print('Firing rate at neuron {}: \n'.format(n_exc_test_print),firing_rate_array[n_exc_test_print])

        head_direction_indices = get_head_direction_indices(directions, firing_rate_array)
        traj.f_add_result('runs.$.head_direction_indices', head_direction_indices,
                          comment='The HDIs of the excitatory population')


        print("Circle" if idx in circular_indices else "Ellipsoid" )
        print("Corr length {:.1f}".format(traj.orientation_map.correlation_length))
        print("Long axis {:.1f}".format(traj.long_axis))
        print("Mean HDI {:.1f}".format(np.mean(head_direction_indices)))
        traj.f_restore_default()
    traj.f_store()