from brian2.core.namespace import get_local_namespace
from brian2 import Synapses, Function
from brian2.input.spikegeneratorgroup import SpikeGeneratorGroup
from brian2lava.utils.const import HARDWARE

MAX_SYNAPSES = 2147483647


def spike_queue(self, source_start, source_end):
    """
    TODO

    Parameters
    ----------
    source_start
        TODO
    source_end
        TODO
    
    Returns
    -------
    `SpikeQueue`
        TODO
    """

    # TODO REMOVE! In brian2lava there is no need for a spike queue.
    # Use the C++ version of the SpikeQueue when available
    try:
        from brian2.synapses.cythonspikequeue import SpikeQueue
        self.logger.diagnostic('Using the C++ SpikeQueue', once=True)
    except ImportError:
        from brian2.synapses.spikequeue import SpikeQueue
        self.logger.diagnostic('Using the Python SpikeQueue', once=True)

    return SpikeQueue(source_start=source_start, source_end=source_end)


def synapses_connect(
        self, synapses, condition=None, i=None, j=None, p=1., n=1,
        skip_if_invalid=False, namespace=None, level=0
    ):
    """
    Connects synapses.

    This method overwrites the ``connect()`` function from Brian. For flexible mode, the only difference
    to the original Brian function is that we add a Lava-specific variable to the registered variables, 
    and at the end of the connection process we add the code object generated by the synapse to the
    initialization queue. 
    For preset mode, there are no differences to the original Brian function.

    Parameters
    ----------
    synapses
        TODO
    condition : optional
        TODO
    i : optional
        TODO
    j : optional
        TODO
    p : optional
        TODO
    n : optional
        TODO
    skip_if_invalid : optional
        TODO
    namespace : optional
        TODO
    level : optional
        TODO

    Notes
    -----
    TODO Maybe there is a better way to do this which is more maintainable.
    """
    if self.mode == 'flexible':
        # First, add the spiking synapses variable to the lava_variables
        self._add_spiking_synapses_vars(synapses)

    # STANDARD BRIAN CODE -------------------------------------------
    # check types
    synapses._verify_connect_argument_types(condition, i, j, n, p)

    synapses._connect_called = True

    # Get namespace information
    if namespace is None:
        namespace = get_local_namespace(level=level + 2)

    try:  # wrap everything to catch IndexError
        # which connection case are we in?
        # 1: Connection condition
        if condition is None and i is None and j is None:
            condition = True
        if condition is not None:
            if i is not None or j is not None:
                raise ValueError("Cannot combine condition with i or j arguments")
            if condition is False or condition == "False":
                # Nothing to do
                return
            j = synapses._condition_to_generator_expression(condition, p, namespace)
            synapses._add_synapses_generator(
                j,
                n,
                skip_if_invalid=skip_if_invalid,
                namespace=namespace,
                level=level + 2,
                over_presynaptic=True,
            )
        # 2: connection indices
        elif (i is not None and j is not None) and not (
            isinstance(i, str) or isinstance(j, str)
        ):
            if skip_if_invalid:
                raise ValueError("Can only use skip_if_invalid with string syntax")
            i, j, n = synapses._verify_connect_array_arguments(i, j, n)
            synapses._add_synapses_from_arrays(i, j, n, p, namespace=namespace)
        # 3: Generator expression over post-synaptic cells (i='...')
        elif isinstance(i, str):
            i = synapses._finalize_generator_expression(i, j, p, "i", "j")
            synapses._add_synapses_generator(
                i,
                n,
                skip_if_invalid=skip_if_invalid,
                namespace=namespace,
                level=level + 2,
                over_presynaptic=False,
            )
        # 4: Generator expression over pre-synaptic cells (i='...')
        elif isinstance(j, str):
            j = synapses._finalize_generator_expression(j, i, p, "j", "i")
            synapses._add_synapses_generator(
                j,
                n,
                skip_if_invalid=skip_if_invalid,
                namespace=namespace,
                level=level + 2,
                over_presynaptic=True,
            )
        else:
            raise ValueError(
                "Must specify at least one of condition, i or j arguments"
            )
    except IndexError as e:
        raise IndexError(
            "Tried to create synapse indices outside valid "
            "range. Original error message: "
            + str(e)
        )
    # ----------------------------------------------------------------------    
    if self.mode == 'flexible':
        # Add the codeobject determining synaptic connections to the init queue for synapses
        generator_objects = [codeobj for codeobj in self.code_objects.values() if codeobj.owner == synapses and 'synapses_create' in codeobj.template_name]

        # Make sure that we don't insert duplicates in the queue
        for func, args in self.proc_init_queue[synapses.name]:
            if func == 'code_object':
                codeobj = args
                # If the codeobject was already added during a potential previous call of connect
                if codeobj in generator_objects:
                    generator_objects.remove(codeobj)
        
        # Insert the new generators in the init queue
        for codeobj in generator_objects:
            self.proc_init_queue[synapses.name].append(('code_object', codeobj))


def _add_spiking_synapses_vars(self, synapses):
    """
    Adds a variable to the list of Lava variables, which will store the spiking synapses from the last
    timestep. The vector is then used to update the synaptic variables in the ``run_lrn`` function if
    the synapses are plastic.

    Parameters
    ----------
    TODO
    """
    for pathway in synapses._pathways:

        # The spiking_synapses variable is always added
        spiking_synapses = f"spiking_{synapses.name}_{pathway.objname}"
        size = 1

        # Only add the variable once for each pathway.
        if not spiking_synapses in synapses.variables.keys():
            # This needs to be a dynamic array as its shape will change during the simulation
            synapses.variables.add_dynamic_array(spiking_synapses,size,dtype = "bool")

            # Register this variable so that it will be correctly resized during the initialization.
            synapses.register_variable(synapses.variables[spiking_synapses])
        

def determine_lava_ports(self, pathway, variables):
    """
    For a specified synyaptic pathway, extracts the synaptic variables to be sent 
    to the neuron groups and determine the Lava ports required for that.

    Parameters
    ----------
    pathway : `SynapticPathway`
        The synaptic pathway object.
    variables : dict of `any`
        Dictionary of the variables contained in the code that are used by the pathway.
    """
    # Support for subgroups will come in the future
    if 'subgroup' in pathway.source.name or 'subgroup' in pathway.target.name:
        msg = f"""
        Encountered a Subgroup object in the connection from '{pathway.source.name}' to '{pathway.target.name}'.
        Subgroups are not supported yet. Subgroups are defined by indexing a NeuronGroup (e.g. 'P = NG[0:10]').
        If you are using this method to connect synapses we suggest using alternative methods such as
        the expression 'S.connect('i<10', p = p)', which has been tested.
        """
        raise NotImplementedError(msg)
    synaptic_vars = []
    # TODO: this is probably not needed in case the update only requires synaptic variables,
    # then you wouldn't need to send them to the neuronGroup. But for now let's just make
    # the easiest scenario
    for varname,var in variables.items():
        if isinstance(var, Function) or not isinstance(var.owner, Synapses):
            continue
        synaptic_vars.append(varname)

    source_name = pathway.source.name
    target_name = pathway.target.name
    
    # Currently in preset mode only the synaptic variable `s` is supported
    if self.mode == 'preset':
        synaptic_vars = ['s']
        # On Loihi, spike generators need an additional process/adapter.
        if isinstance(pathway.source, SpikeGeneratorGroup) and self.hardware == HARDWARE.Loihi2:
            # Add a an intermediate port that connects to the adapter
            portname = pathway.synapses.name + pathway.objname + '_adapter'
            self.lava_ports[portname] = {
                'pathway': pathway,
                'sender': source_name,
                'receiver': source_name + '_adapter'
            }
            # Change the source name so that the intermediate adapter is the new source
            source_name = source_name + '_adapter'

    # Variable names are kept in their original form, but ports are defined separately depending on the pathway.
    for varname in synaptic_vars:
        portname = pathway.synapses.name + '_' + pathway.objname + '_' + varname
        self.lava_ports[portname] = {
            'varname': varname, # Required by code generation
            'portname': portname,
            'pathway': pathway,
            'sender' : source_name,
            'receiver': target_name,
        }

    ports = '\n'.join([port for port in self.lava_ports])
    msg = f"""Stored ports for synaptic transmission from '{source_name}' to '{target_name}' via '{pathway.synapses.name}':
           {ports}
           """
    self.logger.diagnostic(msg)

