import unittest

from connectome.analysis import ConnectivityEstimator
from connectome.model.network import Network


class TextConnectivityEstimator(unittest.TestCase):
    def test_connectivity_estimator_zero_connectivity(self):
        adjacency = [[0,0,0],
                     [0,0,0],
                     [0,0,0]]
        net = Network(adjacency, nr_exc=2, nr_inh=1)
        connectivity_estimator = ConnectivityEstimator()
        for key, value in connectivity_estimator(network=net).items():
            self.assertEqual(0, value)

    def test_connectivity_estimator_inh_only(self):
        adjacency = [[0,0,0,0],
                     [0,0,0,0],
                     [0,0,0,-1],
                     [0,0,-1,0]]
        net = Network(adjacency, nr_exc=2, nr_inh=2)
        connectivity_estimator = ConnectivityEstimator()
        self.assertEqual(1, connectivity_estimator(network=net)['p_ii'])

    def test_connectivity_estimator_exc_only(self):
        adjacency = [[0,.1,0,0],
                     [1,0,0,0],
                     [0,0,0,0],
                     [0,0,0,0]]
        net = Network(adjacency, nr_exc=2, nr_inh=2)
        connectivity_estimator = ConnectivityEstimator()
        self.assertEqual(1, connectivity_estimator(network=net)['p_ee'])

    def test_connectivity_estimator_exc_inh(self):
        adjacency = [[0,0,0,0],
                     [0,0,0,0],
                     [1,.1,0,0],
                     [1,1,0,0]]
        net = Network(adjacency, nr_exc=2, nr_inh=2)
        connectivity_estimator = ConnectivityEstimator()
        self.assertEqual(1, connectivity_estimator(network=net)['p_ei'])

    def test_connectivity_estimator_inh_exc(self):
        adjacency = [[0,0,-1,-1],
                     [0,0,-1,-1],
                     [0,0,0,0],
                     [0,0,0,0]]
        net = Network(adjacency, nr_exc=2, nr_inh=2)
        connectivity_estimator = ConnectivityEstimator()
        self.assertEqual(1, connectivity_estimator(network=net)['p_ie'])

    def test_connectivity_estimator_mixed_ee(self):
        adjacency = [[0,1,-1,0],
                     [0,0,-1,0],
                     [0,0,0,-1],
                     [1,1,0,0]]
        net = Network(adjacency, nr_exc=2, nr_inh=2)
        connectivity_estimator = ConnectivityEstimator()
        self.assertEqual(1/2, connectivity_estimator(network=net)['p_ee'])

    def test_connectivity_estimator_mixed_ii(self):
        adjacency = [[0,1,-1,0],
                     [0,0,-1,0],
                     [0,0,0,-1],
                     [1,1,0,0]]
        net = Network(adjacency, nr_exc=2, nr_inh=2)
        connectivity_estimator = ConnectivityEstimator()
        self.assertEqual(1/2, connectivity_estimator(network=net)['p_ii'])

    def test_connectivity_estimator_mixed_ei(self):
        adjacency = [[0,1,-1,0],
                     [0,0,-1,0],
                     [0,0,0,-1],
                     [1,1,0,0]]
        net = Network(adjacency, nr_exc=2, nr_inh=2)
        connectivity_estimator = ConnectivityEstimator()
        self.assertEqual(2/4, connectivity_estimator(network=net)['p_ei'])

    def test_connectivity_estimator_mixed_ie(self):
        adjacency = [[0,1,-1,0],
                     [0,0,-1,0],
                     [0,0,0,-1],
                     [1,1,0,0]]
        net = Network(adjacency, nr_exc=2, nr_inh=2)
        connectivity_estimator = ConnectivityEstimator()
        self.assertEqual(2/4, connectivity_estimator(network=net)['p_ie'])

    def test_output_filtering(self):
        adjacency = [[0,1,-1,0],
                     [0,0,-1,0],
                     [0,0,0,-1],
                     [1,1,0,0]]
        net = Network(adjacency, nr_exc=2, nr_inh=2)
        measures = ["p_ee", "p_ei"]
        connectivity_estimator = ConnectivityEstimator(measures=measures)
        results = connectivity_estimator(network=net)
        self.assertEqual(set(measures), set(results.keys()))

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