import unittest

from connectome.model import ER
from connectome.noise import OutDegreePreservingNoise


class TestOutDegreePreservingNoiseShufflePartially(unittest.TestCase):
    def setUp(self):
        er = ER(nr_neurons=500, inh_ratio=3/5, p_exc=.4, p_inh=.4)
        unshuffled_network = er()
        out_deg_preserving_noise = OutDegreePreservingNoise(network=unshuffled_network,
                                                     fraction_draws_over_nr_synapses=.3)
        shuffled_network = out_deg_preserving_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.out_deg_preserving_noise = out_deg_preserving_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_out_degrees_preserved(self):
        for original, noisy in self.submatrices:
            self.assertTrue(((original != 0).sum(0) == (noisy != 0).sum(0)).all())  # out degree

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

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


class TestOutDegreePreservingNoiseShuffleFull(unittest.TestCase):
    def setUp(self):
        er = ER(nr_neurons=500, inh_ratio=2/5, p_exc=.4, p_inh=.4)
        unshuffled_network = er()
        out_deg_preserving_noise = OutDegreePreservingNoise(network=unshuffled_network,
                                                     fraction_draws_over_nr_synapses=2)
        shuffled_network = out_deg_preserving_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.out_deg_preserving_noise = out_deg_preserving_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_out_degrees_preserved(self):
        for original, noisy in self.submatrices:
            self.assertTrue(((original != 0).sum(0) == (noisy != 0).sum(0)).all())  # out degree

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

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


class TestOutDegreePreservingNoShuffle(unittest.TestCase):
    def setUp(self):
        er = ER(nr_neurons=600, inh_ratio=.5, p_exc=.4, p_inh=.4)
        unshuffled_network = er()
        out_deg_preserving_noise = OutDegreePreservingNoise(network=unshuffled_network,
                                                     fraction_draws_over_nr_synapses=0)
        shuffled_network = out_deg_preserving_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.out_deg_preserving_noise = out_deg_preserving_noise

    def test_not_modified(self):
        self.assertEqual(self.unshuffled_network, self.shuffled_network)


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