from modelflow import Node, Input
import scipy as sp


class AddEdgesNoise(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_add = 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_add() == 0:
            return noisy_network
        self._add_edges(noisy_network)
        noisy_network.binarize()
        return noisy_network

    def _add_edges(self, noisy_network):
        candidate_targets, candidate_sources = self._free_edge_locations(noisy_network)
        nr_candidates = candidate_targets.size

        new_edges = sp.random.choice(nr_candidates, self._nr_edges_to_add(), replace=False)
        new_targets = candidate_targets[new_edges]
        new_sources = candidate_sources[new_edges]

        new_exc_edges = new_sources < self.network.nr_exc
        new_inh_edges = new_sources >= self.network.nr_exc

        noisy_network.adjacency_matrix[new_targets[new_exc_edges], new_sources[new_exc_edges]] = 1
        noisy_network.adjacency_matrix[new_targets[new_inh_edges], new_sources[new_inh_edges]] = -1

    def _nr_edges_to_add(self):
        return int(sp.count_nonzero(self.network.adjacency_matrix) * self.fraction_add)

    @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
