import brian2
import numpy
from interneuron_polarity.io.io import IO
from interneuron_polarity.model.dynamics.brian_utils import set_parameters_from_dict, convert_string_to_unit


def simulate_network_activity(path_to_simulation_dir):
    my_io = IO(path_to_simulation_dir)

    dynamics_dict = my_io.load_dynamics_dict()
    in_ex_connectivity = my_io.load_in_ex_connectivity()
    experiment_dict = my_io.load_experiment()

    excitatory_neurons, inhibitory_neurons, in_ex_synapses, ex_in_synapses = create_neuronal_populations_and_synapses(
        dynamics_dict,
        in_ex_connectivity)


    # initialization
    set_parameters_from_dict(excitatory_neurons,
                             {"v": experiment_dict["excitatory"]["v_reset"]})
    set_parameters_from_dict(inhibitory_neurons,
                             {"v": experiment_dict["inhibitory"]["v_reset"]})

    excitatory_spike_recorder = brian2.SpikeMonitor(source=excitatory_neurons)
    inhibitory_spike_recorder = brian2.SpikeMonitor(source=inhibitory_neurons)
    recorders = [excitatory_spike_recorder, inhibitory_spike_recorder]
    run_experiment(experiment_dict, excitatory_neurons, inhibitory_neurons, ex_in_synapses, in_ex_synapses,
                   recorders)

    excitatory_spikes = excitatory_spike_recorder.spike_trains()
    inhibitory_spikes = inhibitory_spike_recorder.spike_trains()
    my_io.save_spike_trains(excitatory_spikes, inhibitory_spikes)

    # What is the best way to setup the network?
    # The network has a skeleton that is neuronal dynamics, synaptic dynamics and connectivity
    # The network also has some experimental conditions, that is setting neuronal parameters, synaptic parameters and
    # input
    # The simulation further has settings such as methods and time steps
    # It would be nice to split this in multiple steps.
    # Is the network object enough to access all its objects?


def create_neuronal_populations_and_synapses(dynamics_dict, in_ex_connectivity):
    n_in = in_ex_connectivity.shape[0]
    n_ex = in_ex_connectivity.shape[1]

    if isinstance(dynamics_dict["excitatory"]["model"], list):
        convert_model_list_into_a_str(dynamics_dict)

    excitatory_neurons = brian2.NeuronGroup(n_ex, **dynamics_dict["excitatory"])
    inhibitory_neurons = brian2.NeuronGroup(n_in, **dynamics_dict["inhibitory"])

    ex_in_synapses = brian2.Synapses(source=excitatory_neurons, target=inhibitory_neurons, **dynamics_dict["ex-in"])
    in_ex_synapses = brian2.Synapses(source=inhibitory_neurons, target=excitatory_neurons, **dynamics_dict["in-ex"])

    connected_in_ex_pairs = numpy.vstack(numpy.where(in_ex_connectivity))

    excitatory_indices = connected_in_ex_pairs[1]
    inhibitory_indices = connected_in_ex_pairs[0]

    ex_in_synapses.connect(i=excitatory_indices, j=inhibitory_indices)
    in_ex_synapses.connect(i=inhibitory_indices, j=excitatory_indices)
    return excitatory_neurons, inhibitory_neurons, in_ex_synapses, ex_in_synapses


def convert_model_list_into_a_str(dynamics_dict):
    dynamics_dict["excitatory"].update({"model": "\n".join(dynamics_dict["excitatory"]["model"])})
    dynamics_dict["inhibitory"].update({"model": "\n".join(dynamics_dict["inhibitory"]["model"])})


def run_experiment(experiment_dict, excitatory_neurons, inhibitory_neurons, ex_in_synapses, in_ex_synapses,
                   recorders):
    set_parameters_from_dict(excitatory_neurons, experiment_dict["excitatory"])
    set_parameters_from_dict(inhibitory_neurons, experiment_dict["inhibitory"])
    set_parameters_from_dict(ex_in_synapses, experiment_dict["ex-in"])
    set_parameters_from_dict(in_ex_synapses, experiment_dict["in-ex"])
    net = brian2.Network(excitatory_neurons, inhibitory_neurons)
    net.add((ex_in_synapses, in_ex_synapses))
    net.add(tuple(recorders))
    duration = convert_string_to_unit(experiment_dict["duration"])
    net.run(duration, report="stdout")  # make nice tqdm report
