from modelflow import StaticGraph
from modelflow.elementarynodes import Input
from .subsampling import Subsampling
from .removeandaddedgesnoise import RemoveAndAddEdgesNoise


class RemoveAddNoiseAndSubsample(StaticGraph):
    """
    Subsample a fraction of the network and shuffle edges in the remaining portion.

    Parameters
    ----------

    network: :class:`Network <connectome.model.network.Network>`
        The network.

    fraction_remove_and_add: float in [0, 1]
        Fraction of edges to be removed and randomly reinserted again.
        The sign of an edge is automatically converted to match the type
        of the new presynaptic neuron.

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

    def construct_graph(self):
        self.add_path(({"network" : "network",
                        "fraction_remove_and_add": "fraction_remove_and_add"}, RemoveAndAddEdgesNoise()),
                      ({"%": "network",
                        "subsampling_fraction": "subsampling_fraction"}, Subsampling()),
                      last_is_output=True)
