from .edgeshufflingnoise import EdgeShufflingNoise

# due to maybe not correctly installed gsl libraries
try:
    from connectome.shuffling import shuffle_graph
except ImportError:
    shuffle_graph = None


class InAndOutDegreePreservingNoise(EdgeShufflingNoise):
    """
    Shuffle edges in a way such that in and out degrees are preserved.

    This is a subclass of :class:`EdgeShufflingNoise <connectome.noise.edgeshufflingnoise.EdgeShufflingNoise>`.

    The shuffling is achieved through a Markov-Chain-Monte-Carlo method.
    An ergodic Markov chain on the space of graphs is constructed.
    It is then sampled from this Markov chain.
    The sample is thus an unbiased sample from all graphs with the given in
    and out degree distribution.
    """
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.last_nr_square_swaps = None
        self.last_nr_triangle_swaps = None

    def run(self):
        self.last_nr_square_swaps = 0
        self.last_nr_triangle_swaps = 0
        return super().run()

    def _shuffle_subpopulation(self, submatrix, diagonal):
        nr_draws = int((submatrix != 0).sum() * self.fraction_draws_over_nr_synapses)
        shuffled_submatrix, nr_square_swaps, nr_triangle_swaps = shuffle_graph(submatrix, diagonal, nr_draws, triangle_moves=False)
        self.last_nr_square_swaps += nr_square_swaps
        self.last_nr_triangle_swaps += nr_triangle_swaps
        submatrix[:, :] = shuffled_submatrix
