import numpy as np
from brian2.units import *
from pypet import Environment, cartesian_product, Trajectory
from pypet.brian2.parameter import Brian2MonitorResult

from scripts.interneuron_placement import create_grid_of_excitatory_neurons, \
    create_interneuron_sheet_entropy_max_orientation, get_excitatory_neurons_in_inhibitory_axonal_clouds
from scripts.ring_network.head_direction import ex_in_network
from scripts.spatial_maps.orientation_map import OrientationMap
from scripts.spatial_maps.orientation_map_generator_pypet import TRAJ_NAME_ORIENTATION_MAPS
from scripts.spatial_network.head_direction_index_over_noise_scale import excitatory_eqs, excitatory_params, \
    lif_interneuron_eqs, lif_interneuron_params, lif_interneuron_options, ei_synapse_model, ei_synapse_on_pre, \
    ei_synapse_param, ie_synapse_model, ie_synapse_on_pre, ie_synapse_param, get_synaptic_weights, \
    create_head_direction_input

DATA_FOLDER = "../../data/"
LOG_FOLDER = "../../logs/"

TRAJ_NAME = "entropy_maximisation_over_different_input_correlation_lengths"

n_exc_test_print = 465

def get_orientation_map(correlation_length, seed, sheet_size, N_E):
    traj = Trajectory(filename=DATA_FOLDER + TRAJ_NAME_ORIENTATION_MAPS + ".hdf5")

    traj.f_load(index=-1, load_parameters=2, load_results=2)

    corr_len = correlation_length
    seed = seed

    map_by_params = lambda x, y: x == corr_len and y == seed

    idx_iterator = traj.f_find_idx(['corr_len', 'seed'], map_by_params)

    # TODO: Since it has only one entry, maybe iterator can be replaced
    for idx in idx_iterator:
        traj.v_idx = idx
        map_angle_grid = traj.crun.map

    number_of_excitatory_neurons_per_row = int(np.sqrt(N_E))

    map = OrientationMap(number_of_excitatory_neurons_per_row + 1, number_of_excitatory_neurons_per_row + 1,
                         corr_len, sheet_size, sheet_size, seed)
    map.angle_grid = map_angle_grid

    tuning_map = lambda x, y: map.orientation_map(x, y)
    return tuning_map

# TODO: Rename to avoid namespace conflicts
def experiment(traj):
    sheet_size = traj.orientation_map.sheet_size

    N_E = traj.network.N_E
    N_I = traj.network.N_I

    orientation_map = get_orientation_map(traj.orientation_map.correlation_length, traj.orientation_map.seed,
                                          sheet_size, N_E)
    ex_positions, ex_tunings = create_grid_of_excitatory_neurons(sheet_size,
                                                                 sheet_size,
                                                                 int(np.sqrt(N_E)), orientation_map)

    inhibitory_axon_long_axis = traj.morphology.long_axis
    inhibitory_axon_short_axis = traj.morphology.short_axis

    # TODO: Remove one step for circles as it's unnecessary
    entropy_maximisation_steps = traj.simulation.entropy_maximisation.steps if inhibitory_axon_long_axis != \
                                                                               inhibitory_axon_short_axis else 1

    inhibitory_axonal_clouds, ellipse_single_trial_entropy = create_interneuron_sheet_entropy_max_orientation(
        ex_positions, ex_tunings, N_I, inhibitory_axon_long_axis,
        inhibitory_axon_short_axis, sheet_size,
        sheet_size, trial_orientations=entropy_maximisation_steps)

    ie_connections = get_excitatory_neurons_in_inhibitory_axonal_clouds(ex_positions, inhibitory_axonal_clouds)

    inhibitory_synapse_strength = traj.synapse.inhibitory * nS
    excitatory_synapse_strength = traj.synapse.excitatory * mV

    if inhibitory_synapse_strength != 0.0 * nS and excitatory_synapse_strength != 0.0 * mV \
        and inhibitory_axon_long_axis == inhibitory_axon_short_axis:
        print('circular')
        traj.f_add_derived_parameter("morphology.morph_label", 'circular',
                                     comment="Interneuron morphology of this run is circular")
    elif inhibitory_synapse_strength != 0.0 * nS and excitatory_synapse_strength != 0.0 * mV:
        print('ellipsoid')
        traj.f_add_derived_parameter("morphology.morph_label", 'ellipsoid',
                                     comment="Interneuron morphology of this run is ellipsoid")
    else:
        # TODO: implement non-connected case for comparison
        print('no conn')
        traj.f_add_derived_parameter("morphology.morph_label", 'no conn',
                                     comment="There are no interneurons")

    ex_in_weights, in_ex_weights = get_synaptic_weights(N_E, N_I, ie_connections, excitatory_synapse_strength,
                                                        inhibitory_synapse_strength)

    sharpness = 1.0 / (traj.input.width) ** 2

    # TODO: Parameterize input direction?
    directions = np.linspace(-np.pi, np.pi, traj.input.number_of_directions)
    # TODO: Maybe smarter to use pypet explore instead of loop
    for idx, dir in enumerate(directions):
        # We recreate the network here for every dir, which slows down the simulation quite considerably. Otherwise,
        # we get a problem with saving and restoring the spike times (0s spike for neuron 0)
        net = ex_in_network(N_E, N_I, excitatory_eqs, excitatory_params, lif_interneuron_eqs,
                            lif_interneuron_params,
                            lif_interneuron_options, ei_synapse_model, ei_synapse_on_pre,
                            ei_synapse_param,
                            ex_in_weights, ie_synapse_model, ie_synapse_on_pre,
                            ie_synapse_param, in_ex_weights, random_seed=2)
        input_to_excitatory_population = create_head_direction_input(traj.input.baseline * nA, ex_tunings,
                                                                     sharpness,
                                                                     traj.input.amplitude * nA, dir)

        excitatory_neurons = net["excitatory_neurons"]
        excitatory_neurons.I = input_to_excitatory_population
        net.run(traj.simulation.duration * ms)

        direction_id = 'dir{:d}'.format(idx)
        # print(net["excitatory_spike_monitor"].spike_trains()[n_exc_test_print])
        traj.f_add_result(Brian2MonitorResult, '{:s}.spikes.e'.format(direction_id), net["excitatory_spike_monitor"],
                          comment='The spiketimes of the excitatory population')
        traj.f_add_result(Brian2MonitorResult, '{:s}.spikes.i'.format(direction_id), net["inhibitory_spike_monitor"],
                          comment='The spiketimes of the inhibitory population')
    traj.f_add_result('ex_positions', np.array(ex_positions),
                      comment='The positioons of the excitatory neurons on the sheet')
    traj.f_add_result('ex_tunings', np.array(ex_tunings),
                      comment='The input tunings of the excitatory neurons')
    # TODO: Switch to pandas? Nested lists dont work unfortunately
    print(len(ie_connections))
    ie_connections_save_array = np.zeros((N_I,N_E))
    for i_idx, ie_conn in enumerate(ie_connections):
        for e_idx in ie_conn:
            ie_connections_save_array[i_idx,e_idx] = 1
    traj.f_add_result('ie_adjacency', ie_connections_save_array,
                      comment='Recurrent connection adjacency matrix')

    # TODO: Somewhat ugly, since the objects can not be saved
    axon_cloud_save_list = [[p.x,p.y,p.phi] for p in inhibitory_axonal_clouds]
    axon_cloud_save_array = np.array(axon_cloud_save_list)
    traj.f_add_result('inhibitory_axonal_cloud_array', axon_cloud_save_array,
                      comment='The inhibitory axonal clouds')

    return 1


if __name__ == "__main__":
    env = Environment(trajectory=TRAJ_NAME,
                      comment="Compare the head direction tuning for circular and ellipsoid interneuron morphology, "
                              "when tuning orientations to maximise entropy of connected excitatory tunings.",
                      multiproc=True, filename=DATA_FOLDER, ncores=3, overwrite_file=True, log_folder=LOG_FOLDER)

    traj = env.trajectory

    traj.f_add_parameter_group("orientation_map")

    traj.f_add_parameter("orientation_map.correlation_length", 200.0, comment="Correlation length of orientations in um")
    traj.f_add_parameter("orientation_map.seed", 1, comment="Random seed for map generation.")
    traj.f_add_parameter("orientation_map.sheet_size", 450, comment="Sheet size in um")

    traj.f_add_parameter_group("network")
    traj.f_add_parameter("network.N_E", 900, comment="Number of excitatory neurons")
    traj.f_add_parameter("network.N_I", 100, comment="Number of inhibitory neurons")

    traj.f_add_parameter_group("synapse")
    traj.f_add_parameter("synapse.inhibitory", 30, "Strength of conductance-based inhibitory synapse in nS.")
    traj.f_add_parameter("synapse.excitatory", 1, "Strength of conductance-based inhibitory synapse in mV.")

    traj.f_add_parameter_group("input")
    traj.f_add_parameter("input.width", np.pi / 3.0, comment="Standard deviation of incoming head direction input.")
    traj.f_add_parameter("input.baseline", 0.2, comment="Head direction input baseline")
    traj.f_add_parameter("input.amplitude", 0.5, comment="Head direction input amplitude")
    traj.f_add_parameter("input.number_of_directions", 12, comment="Number of probed directions")

    traj.f_add_parameter_group("morphology")
    traj.f_add_parameter("morphology.long_axis", 200.0, comment="Long axis of axon ellipsoid")  # Should
    # this rather
    # be a
    # parameter, so that it is part of the comparison?
    traj.f_add_parameter("morphology.short_axis", 50.0, comment="Short axis of axon ellipsoid")

    traj.f_add_parameter_group("simulation")
    traj.f_add_parameter("simulation.entropy_maximisation.steps", 30, comment="Steps for entropy maximisation")
    traj.f_add_parameter("simulation.dt", 0.1, comment="Network simulation time step in ms")
    traj.f_add_parameter("simulation.duration", 1000, comment="Network simulation duration in ms")

    # TODO: The Orientation Maps need to be precalculated for these values
    # TODO: Random seeds need to be added to exploration
    ellipsoid_parameter_exploration = {
        "morphology.long_axis": [200.0],
        "morphology.short_axis": [50.0],
        "orientation_map.correlation_length": [200.0]
        # "orientation_map.correlation_length": np.arange(0.0, 200.0, 50).tolist()
    }

    corresponding_circular_radius = float(np.sqrt(ellipsoid_parameter_exploration[
                                                      "morphology.long_axis"][0] * ellipsoid_parameter_exploration[
                                                      "morphology.short_axis"][0]))

    circle_parameter_exploration = {
        "morphology.long_axis": [corresponding_circular_radius],
        "morphology.short_axis": [corresponding_circular_radius],
        "orientation_map.correlation_length": ellipsoid_parameter_exploration["orientation_map.correlation_length"]
    }

    expanded_dicts = [cartesian_product(dict) for dict in [ellipsoid_parameter_exploration,
                                                           circle_parameter_exploration]]
    final_dict = {}

    for key in ellipsoid_parameter_exploration.keys():
        final_dict[key] = [value for l in list(zip(expanded_dicts[0][key], expanded_dicts[1][key])) for value in l]

    traj.f_explore(final_dict)

    env.run(experiment)

    env.disable_logging()
