import copy
import warnings

import brian2 as br
from brian2 import ms


class TimedInhibition:
    def __init__(self, neuron_model, synapse_model, synapse_on_pre, neuron_parameters, synapse_parameters,
                 initial_states, spike_threshold, integration_method, current_drive, record_variables):
        threshold_eqs = """
        v_threshold: volt
        """

        self.neuron = br.NeuronGroup(N=1, \
                                     model=neuron_model + threshold_eqs, \
                                     threshold='v > v_threshold', \
                                     refractory='v > v_threshold', \
                                     method=integration_method)

        self.neuron.v_threshold = spike_threshold
        set_parameters_from_dict(self.neuron, initial_states)
        self.neuron.I = current_drive

        self.spike_recorder = br.SpikeMonitor(source=self.neuron)
        self.neuron_state_recorder = br.StateMonitor(self.neuron, record_variables, record=True)

        self.auto_synapse = br.Synapses(source=self.neuron, target=self.neuron, model=synapse_model,
                                        on_pre=synapse_on_pre,
                                        delay=0.0 * ms)
        self.auto_synapse.connect()
        set_parameters_from_dict(self.auto_synapse, synapse_parameters)
        self.net = br.Network(self.neuron)
        self.net.add(self.auto_synapse)
        self.net.add(self.spike_recorder)
        self.net.add(self.neuron_state_recorder)
        self.net.store()

        self.network_params = copy.deepcopy(neuron_parameters)
        self.network_params.update(synapse_parameters)

    def run_sim(self, delay=0.0, pulse_strength=0.0, record_states=False, duration=50 * ms):
        self.net.restore()
        self.neuron_state_recorder.record = record_states
        self.auto_synapse.delay = delay
        self.auto_synapse.synaptic_strength = pulse_strength

        self.net.run(duration=duration, namespace=self.network_params)

        return self.spike_recorder, self.neuron_state_recorder


def set_parameters_from_dict(neurongroup, dictionary_of_parameters):
    for param_key, param_value in dictionary_of_parameters.items():
        try:
            neurongroup.__setattr__(param_key, param_value)
        except AttributeError as err:
            warnings.warn("{:s} has no parameter {:s}".format(neurongroup.name, param_key))