import unittest

from connectome.model import ER
from connectome.noise import InAndOutDegreePreservingNoise


class TestInOutDegreePreservingNoise(unittest.TestCase):
    def setUp(self):
        er = ER(nr_neurons=2000, inh_ratio=.5, p_exc=.4, p_inh=.4)
        unshuffled_network = er()
        io_deg_pres_noise = InAndOutDegreePreservingNoise(network=unshuffled_network,
                                                          fraction_draws_over_nr_synapses=1)
        shuffled_network = io_deg_pres_noise()
        submatrices = [(unshuffled_network["E", "E"], shuffled_network["E", "E"]),
                       (unshuffled_network["I", "I"], shuffled_network["I", "I"]),
                       (unshuffled_network["I", "E"], shuffled_network["I", "E"]),
                       (unshuffled_network["E", "I"], shuffled_network["E", "I"])]
        self.submatrices = submatrices
        self.shuffled_network = shuffled_network
        self.unshuffled_network = unshuffled_network
        self.io_deg_pres_noise = io_deg_pres_noise

    def test_modified(self):
        self.assertNotEqual(self.unshuffled_network, self.shuffled_network)

    def test_individual_columns_modified(self):
        for original, noisy in self.submatrices:
            for k in range(original.shape[1]):
                self.assertFalse((original[:,k] == noisy[:,k]).all())

    def test_individual_rows_modified(self):
        for original, noisy in self.submatrices:
            for k in range(original.shape[0]):
                self.assertFalse((original[k,:] == noisy[k,:]).all())

    def test_in_out_degrees_preserved(self):
        for original, noisy in self.submatrices:
            self.assertTrue(((original != 0).sum(1) == (noisy != 0).sum(1)).all())  # in degree
            self.assertTrue(((original != 0).sum(0) == (noisy != 0).sum(0)).all())  # out degree

    def test_diagonal_zero(self):
        self.assertFalse(self.shuffled_network.adjacency_matrix.diagonal().any())

    def test_no_triangle_swaps(self):
        self.assertEqual(0, self.io_deg_pres_noise.last_nr_triangle_swaps)

    def test_did_square_swaps(self):
        self.assertGreater(self.io_deg_pres_noise.last_nr_square_swaps, 0)

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