from modelflow import Node, Input
import scipy as sp


class Subsampling(Node):
    """
    Parameters
    ----------

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

    subsampling_fraction: float in [0, 1]
        Fraction of neurons to keep.
    """
    subsampling_fraction = Input()
    network = Input()

    def run(self):
        noisy_network = self.network.copy()
        nr_neurons_to_remove = int((1-self.subsampling_fraction) * noisy_network.nr_neurons)
        if nr_neurons_to_remove > 0:
            neurons_to_remove = sp.random.choice(noisy_network.nr_neurons,
                                                 nr_neurons_to_remove,
                                                 replace=False)
            noisy_network.remove_neurons(neurons_to_remove)
        return noisy_network
