from modelflow import Node, Input
import scipy as sp


class RemoveEdgesNoise(Node):
    """
    Remove and add edges uniformly.

    Parameters
    ----------

    network: :class:`Network <connectome.model.network.Network>`
        The network to which the noise is applied.

    fraction_remove_and_add: float in [0, 1]
        The fraction of edges to be removed and randomly reinserted.
    """
    network = Input()
    fraction_remove = Input()

    def run(self):
        noisy_network = self.network.copy()
        if noisy_network.adjacency_matrix.diagonal().any():
            raise Exception("No diagonal entries allowed.")
        if self._nr_edges_to_remove() == 0:
            return noisy_network
        self._remove_edges(noisy_network)
        noisy_network.binarize()
        print("REMOVE", self.fraction_remove)
        return noisy_network

    def _remove_edges(self, noisy_network):
        """
        Diagonal (self-edges) can be remove if there are any.
        """
        targets, sources = sp.nonzero(noisy_network.adjacency_matrix)
        nr_edges = targets.size
        edges_to_remove = sp.random.choice(nr_edges, self._nr_edges_to_remove(), replace=False)
        noisy_network.adjacency_matrix[targets[edges_to_remove],
                                       sources[edges_to_remove]] = 0

    def _nr_edges_to_remove(self):
        return int(sp.count_nonzero(self.network.adjacency_matrix) * self.fraction_remove)

    @staticmethod
    def _free_edge_locations(network):
        """
        Diagonal not allowed as new edge location
        """
        possible_locations_matrix = (network.adjacency_matrix == 0).astype(int)
        diagonal = possible_locations_matrix.diagonal()
        diagonal.flags.writeable = True
        diagonal[:] = 0
        possible_new_edge_locations_post, possible_new_edge_locations_pre = sp.nonzero(possible_locations_matrix)
        return possible_new_edge_locations_post, possible_new_edge_locations_pre
