import unittest

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


class TestRelativeReciprocityEstimator(unittest.TestCase):
    def test_all_zero_reciprocity(self):
        adjacency = [[0, 0, 0,-0,-0,-0],
                     [0, 0, 0,-0,-0,-0],
                     [0, 0, 0,-0,-0,-0],
                     [0, 0, 0,-0,-0,-0],
                     [0, 0, 0,-0,-0,-0],
                     [0, 0, 0,-0,-0,-0]]
        net = Network(adjacency, nr_exc=3, nr_inh=3)
        reciprocity_estimator = RelativeReciprocityEstimator(network=net)
        for key, value in reciprocity_estimator().items():
            self.assertEqual(0, value)


class TestZeroRelativeReciprocity(unittest.TestCase):
    def setUp(self):
        adjacency = [[0, 1, 0,-0,-0,-0],
                     [0, 0, 0,-0,-0,-1],
                     [0, 1, 0,-0,-0,-1],
                     [1, 0, 0,-0,-1,-0],
                     [1, 0, 0,-0,-0,-0],
                     [0, 0, 0,-0,-1,-0]]
        net = Network(adjacency, nr_exc=3, nr_inh=3)
        reciprocity_estimator = RelativeReciprocityEstimator(network=net)
        self.reciprocity = reciprocity_estimator()

    def tearDown(self):
        del self.reciprocity

    def test_zero_reciprocity_ee(self):
        self.assertEqual(0, self.reciprocity['relative_reciprocity_ee'])

    def test_zero_reciprocity_ei(self):
        self.assertEqual(0, self.reciprocity['relative_reciprocity_ei'])

    def test_zero_reciprocity_ie(self):
        self.assertEqual(0, self.reciprocity['relative_reciprocity_ie'])

    def test_zero_reciprocity_ii(self):
        self.assertEqual(0, self.reciprocity['relative_reciprocity_ii'])


class TestMixedRelativeReciprocity(unittest.TestCase):
    def setUp(self):
        adjacency = [[0, 1, 0,-0,-0,-0],
                     [1, 0, 0,-0,-0,-1],
                     [0, 1, 0,-0,-0,-1],
                     [1, 0, 0,-0,-1,-1],
                     [1, 0, 0,-1,-0,-1],
                     [0, 1, 0,-1,-1,-0]]
        net = Network(adjacency, nr_exc=3, nr_inh=3)
        reciprocity_estimator = RelativeReciprocityEstimator(network=net)
        self.reciprocity = reciprocity_estimator()

    def test_mixed_reciprocity_ee(self):
        self.assertEqual((2/3)/(3/6), self.reciprocity['relative_reciprocity_ee'])

    def test_mixed_reciprocity_ei(self):
        self.assertEqual((1/3)/(2/9), self.reciprocity['relative_reciprocity_ei'])

    def test_mixed_reciprocity_ie(self):
        self.assertEqual((1/2)/(3/9), self.reciprocity['relative_reciprocity_ie'])

    def test_mixed_reciprocity_ii(self):
        self.assertEqual((1)/(1), self.reciprocity['relative_reciprocity_ii'])


class TestOneRelativeReciprocity(unittest.TestCase):
    def setUp(self):
        adjacency = [[0, 1, 1,-1,-1,-1],
                     [1, 0, 1,-1,-1,-1],
                     [1, 1, 0,-1,-1,-1],
                     [1, 1, 1, 0,-1,-1],
                     [1, 1, 1,-1, 0,-1],
                     [1, 1, 1,-1,-1, 0]]
        net = Network(adjacency, nr_exc=3, nr_inh=3)
        reciprocity_estimator = RelativeReciprocityEstimator(network=net)
        self.reciprocity = reciprocity_estimator()

    def test_mixed_reciprocity_ee(self):
        self.assertEqual(1, self.reciprocity['relative_reciprocity_ee'])

    def test_mixed_reciprocity_ei(self):
        self.assertEqual(1, self.reciprocity['relative_reciprocity_ei'])

    def test_mixed_reciprocity_ie(self):
        self.assertEqual(1, self.reciprocity['relative_reciprocity_ie'])

    def test_mixed_reciprocity_ii(self):
        self.assertEqual(1, self.reciprocity['relative_reciprocity_ii'])


class TestNotOneForRelativeReciprocityButOneForReciprocity(unittest.TestCase):
    def setUp(self):
        adjacency = [[0, .1, 1,-1,-1,-1],
                     [1, 0, 1,-1,-1,-1],
                     [1, 1, 0,-1,-1,-1],
                     [1, 1, 1, 0,-1, 0],
                     [1, 1, 1,-1, 0,-1],
                     [1, 1, 1, 0,-.1, 0]]
        net = Network(adjacency, nr_exc=3, nr_inh=3)
        reciprocity_estimator = RelativeReciprocityEstimator(network=net)
        self.reciprocity = reciprocity_estimator()

    def test_mixed_reciprocity_ee(self):
        self.assertEqual(1, self.reciprocity['relative_reciprocity_ee'])

    def test_mixed_reciprocity_ei(self):
        self.assertEqual(1, self.reciprocity['relative_reciprocity_ei'])

    def test_mixed_reciprocity_ie(self):
        self.assertEqual(1, self.reciprocity['relative_reciprocity_ie'])

    def test_mixed_reciprocity_ii(self):
        self.assertEqual(1/(4/6), self.reciprocity['relative_reciprocity_ii'])


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