import scipy as sp

from connectome.noise import RemoveAndAddEdgesNoise
from modelflow import Node, Input


class BarrelCutNoise(Node):
    """
    Simulate to cut a barrel into two pieces.
    Also simulate network noise.
    For spatially embedded networks, the embedding is taken into account.

    :param f: float
        Fraction of the network to be removed.
    :param p: float
        Noise applied to the network after cutting.
        The noise type is :class:`RemoveAndAddEdgesNoise <connectome.noise.removeandaddedgesnoise.RemoveAndAddEdgesNoise>`.
        The parameter `p` is its parameter `fraction_remove_and_add`.
    """
    f = Input()
    p = Input()
    network = Input()

    def edge_removal_probability(self):
        """
        Uses the empirical distribution :math:`emp` from [#feldmeyer_1999]_:


        ==============================  ===========
        Number synapses per connection  Probability
        ==============================  ===========
        1                               0
        2                               2/11
        3                               5 /11
        4                               2/11
        5                               2/11
        ==============================  ===========

        The expected value is 3.36.

        Assumes that synapses are uniformly distributed throughout
        the barrel. Therefore, probability of having a synpase on
        the "same" side is :math:`(1-f)`.

        The fraction of edges additionally lost due to the cut is

        .. math::
           p_{cut}(f) = \sum_{n=0}^5 Binom(0 |n, f)emp(n).

        Parameters
        ----------
        f: float
            Fraction of the barrel which is still left.



        .. [#feldmeyer_1999] Feldmeyer, Dirk, Veronica Egger, Joachim Lübke, and Bert Sakmann.
                             “Reliable Synaptic Connections between Pairs of Excitatory Layer
                             4 Neurones within a Single ‘barrel’ of Developing Rat Somatosensory Cortex.”
                             The Journal of Physiology 521, no. 1 (1999):
                             169–90. doi:10.1111/j.1469-7793.1999.00169.x.
        """
        return 2/11*(1-self.f)**2 + 5/11*(1-self.f)**3 + 2/11*(1-self.f)**4 + 2/11*(1-self.f)**5

    def run(self):
        network = self.network.copy()
        # check for spatial embedding of network
        # remove according to first coordinate
        try:
            positions = network.positions
        except AttributeError:
            nr_vertices_to_remove = int(self.f * network.nr_neurons)
            neurons_to_remove = sp.random.choice(network.nr_neurons, nr_vertices_to_remove, replace=False)
        else:
            neurons_to_remove = sp.nonzero(positions[:,0] >= self.f)[0]
        finally:
            network.remove_neurons(neurons_to_remove)

        self.apply_remove_edges_due_to_cut(network)
        network = self.apply_remove_and_add_edges_noise(network)
        return network

    def apply_remove_edges_due_to_cut(self, network):
        W = network.adjacency_matrix
        p = self.edge_removal_probability()
        eps = sp.ones(W.shape)
        eps[sp.rand(*W.shape) < p] = 0
        W *= eps

    def apply_remove_and_add_edges_noise(self, network):
        network = RemoveAndAddEdgesNoise(fraction_remove_and_add=self.p, network=network)()
        return network
