# Source code for brian2.synapses.spikequeue

    """
    The spike queue class stores future synaptic events
    produced by a given presynaptic neuron group (or postsynaptic for backward
    propagation in STDP).
    """

    import bisect

    import numpy as np

    from brian2.utils.arrays import calc_repeats
    from brian2.utils.logger import get_logger

    __all__ = ["SpikeQueue"]

    logger = get_logger(__name__)

    INITIAL_MAXSPIKESPER_DT = 1

    class SpikeQueue:
        """
        Data structure saving the spikes and taking care of delays.

        Parameters
        ----------

        source_start : int
            The start of the source indices (for subgroups)
        source_end : int
            The end of the source indices (for subgroups)
        Notes
        -----
        **Data structure**

        A spike queue is implemented as a 2D array `X` that is circular in the time
        direction (rows) and dynamic in the events direction (columns). The
        row index corresponding to the current timestep is `currentime`.
        Each element contains the target synapse index.

        **Offsets**

        Offsets are used to solve the problem of inserting multiple synaptic events
        with the same delay. This is difficult to vectorise. If there are n synaptic
        events with the same delay, these events are given an offset between 0 and
        n-1, corresponding to their relative position in the data structure.
        """

        def __init__(self, source_start, source_end):
            #: The start of the source indices (for subgroups)
            self._source_start = source_start

            #: The end of the source indices (for subgroups)
            self._source_end = source_end

            self.dtype = np.int32  # TODO: Ths is fixed for now
            self.X = np.zeros((1, 1), dtype=self.dtype)  # target synapses
            self.X_flat = self.X.reshape(
                1,
            )
            #: The current time (in time steps)
            self.currenttime = 0
            #: number of events in each time step
            self.n = np.zeros(1, dtype=int)

            #: The dt used for storing the spikes (will be set in `prepare`)
            self._dt = None

            self._state_tuple = (self._source_start, self._source_end, self.dtype)

        def prepare(self, delays, dt, synapse_sources):
            """
            Prepare the data structures

            This is called every time the network is run. The size of the
            of the data structure (number of rows) is adjusted to fit the maximum
            delay in `delays`, if necessary. A flag is set if delays are
            homogeneous, in which case insertion will use a faster method
            implemented in `insert_homogeneous`.
            """
            if self._dt is not None:
                # store the current spikes
                spikes = self._extract_spikes()
                # adapt the spikes to the new dt if it changed
                if self._dt != dt:
                    spiketimes = spikes[:, 0] * self._dt
                    spikes[:, 0] = np.round(spiketimes / dt).astype(int)
            else:
                spikes = None

            if len(delays):
                delays = np.array(np.round(delays / dt)).astype(int)
                max_delays = max(delays)
                min_delays = min(delays)
            else:
                max_delays = min_delays = 0

            self._delays = delays

            # Prepare the data structure used in propagation
            synapse_sources = synapse_sources[:]
            ss = np.ravel(synapse_sources)
            # mergesort to retain relative order, keeps the output lists in sorted order
            ss_sort_indices = np.argsort(ss, kind="mergesort")
            ss_sorted = ss[ss_sort_indices]
            splitinds = np.searchsorted(
                ss_sorted, np.arange(self._source_start, self._source_end + 1)
            )
            self._neurons_to_synapses = [
                ss_sort_indices[splitinds[j] : splitinds[j + 1]]
                for j in range(len(splitinds) - 1)
            ]
            max_events = max(map(len, self._neurons_to_synapses))

            n_steps = max_delays + 1

            # Adjust the maximum delay and number of events per timestep if necessary
            # Check if delays are homogeneous
            self._homogeneous = max_delays == min_delays

            # Resize
            if (n_steps > self.X.shape[0]) or (max_events > self.X.shape[1]):  # Resize
                # Choose max_delay if is is larger than the maximum delay
                n_steps = max(n_steps, self.X.shape[0])
                max_events = max(max_events, self.X.shape[1])
                self.X = np.zeros(
                    (n_steps, max_events), dtype=self.dtype
                )  # target synapses
                self.X_flat = self.X.reshape(
                    n_steps * max_events,
                )
                self.n = np.zeros(n_steps, dtype=int)  # number of events in each time step

            # Re-insert the spikes into the data structure
            if spikes is not None:
                self._store_spikes(spikes)

            self._dt = dt

        def _extract_spikes(self):
            """
            Get all the stored spikes

            Returns
            -------
            spikes : ndarray
                A 2d array with two columns, where each row describes a spike.
                The first column gives the time (as integer time steps) and the
                second column gives the index of the target synapse.
            """
            spikes = np.zeros((np.sum(self.n), 2), dtype=int)
            counter = 0
            for idx, n in enumerate(self.n):
                t = (idx - self.currenttime) % len(self.n)
                for target in self.X[idx, :n]:
                    spikes[counter, :] = np.array([t, target])
                    counter += 1
            return spikes

        def _store_spikes(self, spikes):
            """
            Store a list of spikes at the given positions after clearing all
            spikes in the queue.

            Parameters
            ----------
            spikes : ndarray
                A 2d array with two columns, where each row describes a spike.
                The first column gives the time (as integer time steps) and the
                second column gives the index of the target synapse.

            """
            # Clear all spikes
            self.n[:] = 0
            for t, target in spikes:
                row_idx = (t + self.currenttime) % len(self.n)
                self.X[row_idx, self.n[row_idx]] = target
                self.n[row_idx] += 1

        def _full_state(self):
            return (self._dt, self._extract_spikes(), self.X.shape)

        def _restore_from_full_state(self, state):
            if state is None:
                # It is possible that _full_state was called in `SynapticPathway`,
                # before the `SpikeQueue` was created. In that case, delete all spikes in
                # the queue
                self._store_spikes(np.empty((0, 2), dtype=int))
                self._dt = None
            else:
                self._dt, spikes, X_shape = state
                # Restore the previous shape
                n_steps, max_events = X_shape
                self.X = np.zeros((n_steps, max_events), dtype=self.dtype)
                self.X_flat = self.X.reshape(
                    n_steps * max_events,
                )
                self.n = np.zeros(n_steps, dtype=int)
                self._store_spikes(spikes)

        ################################ SPIKE QUEUE DATASTRUCTURE ################

        def advance(self):
            """
            Advances by one timestep
            """
            self.n[self.currenttime] = 0  # erase
            self.currenttime = (self.currenttime + 1) % len(self.n)

        def peek(self):
            """
            Returns the all the synaptic events corresponding to the current time,
            as an array of synapse indexes.
            """
            return self.X[self.currenttime, : self.n[self.currenttime]]

        def push(self, sources):
            """
            Push spikes to the queue.

            Parameters
            ----------
            sources : ndarray of int
                The indices of the neurons that spiked.
            """
            if len(sources) and len(self._delays):
                start = self._source_start
                stop = self._source_end
                if start > 0:
                    start_idx = bisect.bisect_left(sources, start)
                else:
                    start_idx = 0
                if stop <= sources[-1]:
                    stop_idx = bisect.bisect_left(sources, stop, lo=start_idx)
                else:
                    stop_idx = len(sources) + 1
                sources = sources[start_idx:stop_idx]
                if len(sources) == 0:
                    return
                synapse_indices = self._neurons_to_synapses
                indices = np.concatenate(
                    [synapse_indices[source - start] for source in sources]
                ).astype(np.int32)
                if self._homogeneous:  # homogeneous delays
                    self._insert_homogeneous(self._delays[0], indices)
                else:  # vectorise over synaptic events
                    self._insert(self._delays[indices], indices)

        def _insert(self, delay, target):
            """
            Vectorised insertion of spike events.

            Parameters
            ----------

            delay : ndarray
                Delays in timesteps.

            target : ndarray
                Target synaptic indices.
            """
            delay = np.array(delay, dtype=int)

            offset = calc_repeats(delay)

            # Calculate row indices in the data structure
            timesteps = (self.currenttime + delay) % len(self.n)
            # (Over)estimate the number of events to be stored, to resize the array
            # It's an overestimation for the current time, but I believe a good one
            # for future events
            m = max(self.n) + len(target)
            if m >= self.X.shape[1]:  # overflow
                self._resize(m + 1)

            self.X_flat[timesteps * self.X.shape[1] + offset + self.n[timesteps]] = target
            self.n[timesteps] += offset + 1  # that's a trick (to update stack size)

        def _insert_homogeneous(self, delay, target):
            """
            Inserts events at a fixed delay.

            Parameters
            ----------
            delay : int
                Delay in timesteps.

            target : ndarray
                Target synaptic indices.
            """
            timestep = (self.currenttime + delay) % len(self.n)
            nevents = len(target)
            m = (
                self.n[timestep] + nevents + 1
            )  # If overflow, then at least one self.n is bigger than the size
            if m >= self.X.shape[1]:
                self._resize(m + 1)  # was m previously (not enough)
            k = timestep * self.X.shape[1] + self.n[timestep]
            self.X_flat[k : k + nevents] = target
            self.n[timestep] += nevents

        def _resize(self, maxevents):
            """
            Resizes the underlying data structure (number of columns = spikes per
            dt).

            Parameters
            ----------
            maxevents : int
                The new number of columns. It will be rounded to the closest power
                of 2.
            """
            # old and new sizes
            old_maxevents = self.X.shape[1]
            new_maxevents = int(2 ** np.ceil(np.log2(maxevents)))  # maybe 2 is too large
            # new array
            newX = np.zeros((self.X.shape[0], new_maxevents), dtype=self.X.dtype)
            newX[:, :old_maxevents] = self.X[:, :old_maxevents]  # copy old data

            self.X = newX
            self.X_flat = self.X.reshape(
                self.X.shape[0] * new_maxevents,
            )