import unittest

import scipy as sp

from connectome.model import Network
from connectome.noise import RemoveAddNoiseAndSubsample


class TestRemoveAddAndSubsample(unittest.TestCase):
    def test_nr_neurons(self):
        net = Network(nr_exc=50, nr_inh=50)
        remove_and_subsample = RemoveAddNoiseAndSubsample(fraction_remove_and_add =.5,
                                                          subsampling_fraction= .5)
        noisy_net = remove_and_subsample(network=net)
        self.assertEqual(50, noisy_net.nr_neurons)

    def test_shuffling_full_network(self):
        ad = sp.rand(100,100)
        ad[:,25:] *= -1
        diag = ad.diagonal()
        diag.flags.writeable = True
        diag[:] = 0
        net = Network(adjacency_matrix=ad.copy(), nr_exc=25, nr_inh=25)
        remove_and_subsample = RemoveAddNoiseAndSubsample(fraction_remove_and_add =.5,
                                                          subsampling_fraction= .5)
        noisy_net = remove_and_subsample(network=net)
        self.assertNotEqual(ad, noisy_net.adjacency_matrix)


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