import unittest

import numpy as np

from connectome.model import ER, Network
from connectome.noise import RemoveAndAddEdgesNoise


class TestRemoveAndAddFullNetwork(unittest.TestCase):
    def test_full_network(self):
        adj = np.ones((50,50))
        adj[:,50:] *= -1
        diag = adj.diagonal()
        diag.flags.writeable = True
        diag[:] = 0
        net = Network(adj, nr_exc=50, nr_inh=0)
        RemoveAndAddEdgesNoise(network=net, fraction_remove_and_add=.7)()


class TestRemoveAndAddEdgesNoise(unittest.TestCase):
    def setUp(self):
        self.original_network = ER(nr_neurons=100, inh_ratio=0.5, p_exc=.7, p_inh=.7)()
        self.original_network_copy = self.original_network.copy()
        self.noisy_network = RemoveAndAddEdgesNoise(network=self.original_network, fraction_remove_and_add=.7)()

    def test_input_network_unchanged(self):
        self.assertEqual(self.original_network, self.original_network_copy)

    def test_noisy_network_modified(self):
        self.assertNotEqual(self.original_network, self.noisy_network)

    def test_total_nr_edges_preserved(self):
        self.assertEqual(np.count_nonzero(self.original_network.adjacency_matrix),
                         np.count_nonzero(self.noisy_network.adjacency_matrix))

    def test_exc_exc_edges_all_positive(self):
        self.assertTrue((self.noisy_network["E", "E"] >= 0).all())

    def test_exc_inh_edges_all_positive(self):
        self.assertTrue((self.noisy_network["E", "I"] >= 0).all())

    def test_inh_exc_edges_all_negative(self):
        self.assertTrue((self.noisy_network["I", "E"] <= 0).all())

    def test_inh_inh_edges_all_negative(self):
        self.assertTrue((self.noisy_network["I", "I"] <= 0).all())

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


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