import unittest

from connectome.analysis import ConnectivityEstimator
from connectome.model import SYN


class TestSynfire(unittest.TestCase):
    def test_exc_connectivity(self):
        s = SYN(nr_neurons=1000, inh_ratio=.2, pool_size=100, p_inh=.5)
        for p_exc_desired in [.1, .2, .3]:
            net = s(p_exc=p_exc_desired)
            p_exc_obtained = ConnectivityEstimator(network=net)()["p_ee"]
            self.assertLess(abs(1-p_exc_obtained/p_exc_desired), 0.05)

    def test_all_to_all_connectivity(self):
        pool_size = 10
        syn = SYN(nr_neurons=100, inh_ratio=.1, p_exc=.2, p_inh=.5, pool_size=pool_size)
        net = syn()
        adj = net.adjacency_matrix
        pools = syn.pool_connector_.pools_exc

        self.assertGreater(len(pools), 0)

        for k in range(len(pools) - 1):
            pre_pool = pools[k]
            post_pool = pools[k+1]
            for pre_neuron in pre_pool:
                self.assertLess(pre_neuron, net.nr_exc)
                for post_neuron in post_pool:
                    self.assertLess(post_neuron, net.nr_exc)
                    weight = adj[post_neuron, pre_neuron]
                    if pre_neuron != post_neuron:
                        self.assertGreater(weight, 0, msg="Missing connection {}->{}".format(pre_neuron, post_neuron))
                    else:
                        self.assertEqual(weight, 0)


class TestRecurrentSynfireChain(unittest.TestCase):
    def setUp(self):
        self.p_exc = .2
        self.p_inh = .5
        s = SYN(nr_neurons=2000, inh_ratio=.25, pool_size=100, p_inh=self.p_inh, p_exc=self.p_exc)
        self.network = s()

    def test_connectivity_ee(self):
        p_ee = ConnectivityEstimator(network=self.network)()["p_ee"]
        self.assertLess(abs(p_ee / self.p_exc - 1), .05)

    def test_connectivity_ei(self):
        p_ei = ConnectivityEstimator(network=self.network)()["p_ei"]
        # here I allow a rather big error due to the integer rounding
        # of the pool sizes
        self.assertLess(abs(p_ei / self.p_exc - 1), .05)

    def test_connectivity_ee_equals_ei(self):
        p_ee = ConnectivityEstimator(network=self.network)()["p_ee"]
        p_ei = ConnectivityEstimator(network=self.network)()["p_ei"]
        # here I allow a rather big error due to the integer rounding
        # of the pool sizes
        self.assertLess(abs(p_ei / p_ee - 1), .05)

    def test_connectivity_i(self):
        p_i = ConnectivityEstimator(network=self.network)()["p_ii"]
        self.assertLess(abs(p_i / self.p_inh - 1), .01)

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


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