import scipy as sp
from modelflow import Node, Input
from connectome.model.network import Network, EmbeddedNetwork
from typing import Union


class CutFromMiddleNoise(Node):
    """
    Decay power 0.1: about 65% of remaining connections retained (default)
    Decay power 0.8: about 20% retained
    """
    network = Input()
    retained_volume_fraction = Input()
    decay_power = Input()

    SPATIAL_DIMENSION = 3

    def run(self):
        network_cp: Union[Network,EmbeddedNetwork] = self.network.copy()

        try:
            positions = network_cp.positions
        except AttributeError:
            positions = sp.rand(network_cp.nr_neurons, self.SPATIAL_DIMENSION)

        linear_dimension = self.retained_volume_fraction ** (
                    1 / self.SPATIAL_DIMENSION)

        # 0.5 is the barrel center
        onion_layer = (
                    sp.absolute(positions - .5) / (linear_dimension / 2)).max(
            axis=1)
        good_indices = onion_layer <= 1
        good_adjacency = network_cp.adjacency_matrix[good_indices].T[
            good_indices].T
        good_positions = positions[good_indices]
        good_onion_layer = onion_layer[good_indices]

        p_keep_neuron = (1 - good_onion_layer) ** self.decay_power
        p_keep_connection = p_keep_neuron[None].T @ p_keep_neuron[None]
        keep_connection = sp.rand(*p_keep_connection.shape) < p_keep_connection
        good_adjacency[~keep_connection] = 0

        noisy_network = EmbeddedNetwork(good_adjacency)
        noisy_network.positions = good_positions

        return noisy_network
