'''
Written by
  Alessandro Motta <alessandro.motta@brain.mpg.de>
'''

import numpy as np
import os

from connectome.model.network import Network
from modelflow import Node, Input
from typing import Union

class BarrelSubvolumeSimple(Node):
    network = Input()
    cube_length_um = Input()

    SPATIAL_DIMENSION = 3
    BARREL_SIZE_UM = 300

    def run(self):
        keep_frac = (self.cube_length_um / self.BARREL_SIZE_UM) ** self.SPATIAL_DIMENSION

        network_cp: Network = self.network.copy()
        connectome = network_cp.adjacency_matrix

        neuron_count = connectome.shape[0]
        assert all(s == neuron_count for s in connectome.shape)

        # Restrict to neurons within cube
        neuron_mask = np.random.random_sample(neuron_count) < keep_frac
        connectome = connectome[neuron_mask].T[neuron_mask].T

        # Restrict to connections within cube
        conn_ids = np.nonzero(connectome)
        conn_post_ids, conn_pre_ids = conn_ids
        conn_ids = np.ravel_multi_index(conn_ids, connectome.shape)

        conn_count = conn_ids.size
        conn_mask = np.random.random_sample(conn_count) < keep_frac

        connectome_shape = connectome.shape
        connectome = np.reshape(connectome, -1)
        connectome[conn_ids[np.logical_not(conn_mask)]] = 0
        connectome = np.reshape(connectome, connectome_shape)

        noisy_network = Network(connectome)
        return noisy_network
