import numpy as np
import brian2 as br
from brian2.units import *


def ex_in_network(N_E, N_I, excitatory_eqs, excitatory_params, interneuron_eqs, interneuron_params, 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, dt=None):
    # TODO: Check if this does anything
    if random_seed is not None:
        np.random.seed(random_seed)

    if dt is None:
        dt = 0.1 * br.ms
    br.defaultclock = dt

    '''
    Excitatory neurons
    '''
    excitatory_neurons = br.NeuronGroup(N_E, excitatory_eqs, namespace=excitatory_params,
                                        threshold='v>0*mV', refractory='v>0*mV', method='exponential_euler',
                                        name="excitatory_neurons")
    excitatory_neurons.v = excitatory_params["El"] * np.random.normal(1, 0.1, N_E)
    excitatory_spike_monitor = br.SpikeMonitor(excitatory_neurons, name="excitatory_spike_monitor")
    recorded_excitatory_neuron_idx = int(N_E / 2)
    excitatory_trace_recorder = br.StateMonitor(excitatory_neurons, 'v', record=[recorded_excitatory_neuron_idx],
                                                name="excitatory_trace_recorder")
    '''
    Inhibitory neurons
    '''
    inhibitory_neurons = br.NeuronGroup(N_I, interneuron_eqs, namespace=interneuron_params, name="interneurons",
                                        **interneuron_options)
    inhibitory_neurons.v = interneuron_params["u_ext_const"]
    recorded_interneuron_idx = int(N_I / 2)
    inhibitory_trace_recorder = br.StateMonitor(inhibitory_neurons, 'v', record=[recorded_interneuron_idx],
                                                name="inhibitory_trace_recorder")
    inhibitory_spike_monitor = br.SpikeMonitor(inhibitory_neurons, name="inhibitory_spike_monitor")
    '''
    Connectivities
    '''
    ei_synapses = br.Synapses(source=excitatory_neurons, target=inhibitory_neurons, model=ei_synapse_model,
                              on_pre=ei_synapse_on_pre, namespace=ei_synapse_param, name="ei_synapses")
    ei_synapses.connect()

    ex_indices, in_indices = ex_in_weights.nonzero()

    for ex_idx, in_idx in zip(ex_indices, in_indices):
        ei_synapses.synaptic_strength[ex_idx, in_idx] = ex_in_weights[ex_idx, in_idx]
    ie_synapses = br.Synapses(source=inhibitory_neurons, target=excitatory_neurons, model=ie_synapse_model,
                              on_pre=ie_synapse_on_pre, namespace=ie_synapse_param, name="ie_synapses")
    ie_synapses.connect()
    in_indices, ex_indices = in_ex_weights.nonzero()
    for in_idx, ex_idx in zip(in_indices, ex_indices):
        ie_synapses.synaptic_strength[in_idx, ex_idx] = in_ex_weights[in_idx, ex_idx]
    '''
    Net
    '''
    net = br.Network(excitatory_neurons)
    net.add(inhibitory_neurons)
    net.add(ei_synapses)
    net.add(ie_synapses)
    net.add(excitatory_spike_monitor)
    net.add(excitatory_trace_recorder)
    net.add(inhibitory_spike_monitor)
    net.add(inhibitory_trace_recorder)
    net.store()
    return net


def get_head_direction_input(peak_phase, phase_dispersion):
    """
    Returns function that provides drive depending on an array of preferred direction assuming that the input is
    distributed as a
    von
    Mises (
    Gaussian on a ring) distribution with a peak phase and an uncertainty given by the phase dispersion. Normalized
    to one at the peak phase.
    :param peak_phase:
    :param phase_dispersion:
    :return: head_direction_input
    """

    def head_direction_input(direction_preference):
        head_direction_input_shape = np.exp(
            phase_dispersion * np.cos(direction_preference - peak_phase)) / np.exp(phase_dispersion)
        return head_direction_input_shape

    return head_direction_input


def create_head_direction_input(ex_input_baseline, ex_tunings, phase_dispersion, max_head_direction_input_amplitude,
                                tuning_center):
    peak_phase = tuning_center
    direction_input = get_head_direction_input(peak_phase, phase_dispersion)
    input_to_excitatory_population = ex_input_baseline + max_head_direction_input_amplitude * direction_input(
        np.array(ex_tunings))
    return input_to_excitatory_population


def get_synaptic_weights(N_E, N_I, ie_connections, excitatory_synapse_strength, inhibitory_synapse_strength):
    in_ex_weights = np.zeros((N_I, N_E)) * nS
    for interneuron_idx, connected_excitatory_idxs in enumerate(ie_connections):
        in_ex_weights[interneuron_idx, connected_excitatory_idxs] = inhibitory_synapse_strength
    ex_in_weights = np.where(in_ex_weights > 0 * nS, excitatory_synapse_strength, 0 * mV).T * volt
    return ex_in_weights, in_ex_weights