import unittest

from connectome.model import ER, EXP
from connectome.noise import BarrelCutNoise


class NoiseTest(unittest.TestCase):
    def test_barrel_cut_noise_non_spatially_embedded_network(self):
        er = ER(nr_neurons=2000, inh_ratio=.2, p_exc=.2, p_inh=.5)
        network = er()
        barrel_cut_noise = BarrelCutNoise(f=.5, p=0)
        noisy_net = barrel_cut_noise(network=network)
        self.assertLess(noisy_net.nr_neurons, network.nr_neurons // 2 + 1)

    def test_barrel_cut_noise_direct_spatially_embedded_network(self):
        er = EXP(nr_neurons=2000, inh_ratio=.2, p_exc=.2, p_inh=.5, decay=1)
        network = er()
        barrel_cut_noise = BarrelCutNoise(f=.5, p=0)
        noisy_net = barrel_cut_noise(network=network)
        # only check if some are removed.
        # every time a different nr of neurons is removed
        self.assertLess(noisy_net.nr_neurons, network.nr_neurons)
        self.assertFalse((noisy_net.positions[:,0] >= .5).any())


if __name__ == "__main__":
    unittest.main()
