import brian2 as br
import numpy as np
from brian2 import seed


def generate_ring_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,
                          ei_weight_function, ie_synapse_model, ie_synapse_on_pre, ie_synapse_param,
                          ie_weight_function, random_seed=None):

    ex_angles = np.arange(-np.pi, np.pi, 2 * np.pi / N_E)
    in_angles = np.arange(-np.pi, np.pi, 2 * np.pi / N_I)

    ex_in_differences = np.tile(ex_angles, (N_I, 1)).T - np.tile(in_angles, (N_E, 1))
    ex_in_weights = ei_weight_function(ex_in_differences)

    in_ex_difference = ex_in_differences.T

    in_ex_weights = ie_weight_function(in_ex_difference)

    net = 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)

    return ex_angles, in_angles, net


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):

    if random_seed is not None:
        seed(random_seed)
    '''
    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
    '''
    ex_indices, in_indices = ex_in_weights.nonzero()
    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_phase_difference(total_difference):
    """
    Map accumulated phase difference to shortest possible difference
    :param total_difference:
    :return: relative_difference
    """
    return (total_difference + np.pi) % (2 * np.pi) - np.pi


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 get_half_width(x,y):
    """
    Provides the half width of a function with values y given input x
    :param x:
    :param y:
    :return:
    """

    if np.max(y) == np.min(y):
        return 0

    idx_above_half_amplitude = np.asarray(y > np.max(
    y) / 2.0).nonzero()[0]
    half_width_angle_left = x[idx_above_half_amplitude[0]]
    half_width_angle_right = x[idx_above_half_amplitude[-1]]

    return half_width_angle_right-half_width_angle_left