from .edgeshufflingnoise import EdgeShufflingNoise
import scipy as sp


class OutDegreePreservingNoise(EdgeShufflingNoise):
    """
    Shuffle a network such that all out degrees are preserved.

    This is a subclass of :class:`EdgeShufflingNoise <connectome.noise.edgeshufflingnoise.EdgeShufflingNoise>`.
    """
    def _shuffle_subpopulation(self, submatrix, diagonal):
        if self.fraction_draws_over_nr_synapses >= 1:
            self._shuffle_full_matrix(submatrix, diagonal)
        else:
            self._shuffle_partially(submatrix, diagonal)

    def _shuffle_partially(self, submatrix, diagonal):
        nr_edges = sp.count_nonzero(submatrix)
        nr_edges_to_shuffle = int(self.fraction_draws_over_nr_synapses * nr_edges)
        if nr_edges_to_shuffle == 0:
            return

        targets, sources = sp.nonzero(submatrix)
        edges_list = sp.concatenate((sources[None].T, targets[None].T), axis=1)
        edge_ids_to_shuffle = sp.random.choice(nr_edges, nr_edges_to_shuffle, replace=False)
        edges_to_shuffle = edges_list[edge_ids_to_shuffle]
        for neuron_nr in range(submatrix.shape[1]):
            self._shuffle_outgoing_edges_of_single_neuron(neuron_nr, submatrix, edges_to_shuffle, diagonal)

    @staticmethod
    def _shuffle_outgoing_edges_of_single_neuron(neuron_nr, submatrix, edges_to_shuffle, diagonal):
        # edges ends of edges starting at k
        old_targets = edges_to_shuffle[sp.nonzero(edges_to_shuffle[:, 0] == neuron_nr)[0], 1]
        # only do something if there is something to do for neuron  k
        if old_targets.size == 0:
            return
        # store old weights
        weights = submatrix[old_targets, neuron_nr].copy()
        # set old edges to zero
        submatrix[old_targets, neuron_nr] = 0
        # create list of new target candidates
        new_candidates = sp.nonzero(submatrix[:, neuron_nr] == 0)[0]
        # no self connections
        if diagonal:
            new_candidates = new_candidates[new_candidates != neuron_nr]
        new_targets = sp.random.choice(new_candidates, old_targets.size, replace=False)
        # connect
        submatrix[new_targets, neuron_nr] = weights

    @staticmethod
    def _shuffle_full_matrix(submatrix, diagonal):
        """
        For speed. If we shuffle all edges anyway,
        we can speed shuffling up.
        """
        for k in range(submatrix.shape[1]):
            if diagonal:
                tmp = sp.concatenate((submatrix[:k, k], submatrix[k + 1:, k]))
                sp.random.shuffle(tmp)
                submatrix[:k, k] = tmp[:k]
                submatrix[k + 1:, k] = tmp[k:]
            else:
                sp.random.shuffle(submatrix[:, k])