import unittest

from connectome.analysis import ConnectivityEstimator
from connectome.model import API


class TestSimpleMiller(unittest.TestCase):
    def setUp(self):
        self.p_exc = .2
        self.p_inh = .5
        self.simple_miller = API(nr_neurons=2000, inh_ratio=.2, p_exc=self.p_exc,
                                 p_inh=self.p_inh, n_pow=4, feature_space_dimension=30)
        self.net = self.simple_miller()
        self.conn = ConnectivityEstimator(network=self.net)()

    def test_large_population(self):
        self.assertLess(abs(self.conn["p_ee"] / self.p_exc - 1), .01)
        self.assertLess(abs(self.conn["p_ei"] / self.p_exc - 1), .01)
        self.assertLess(abs(self.conn["p_ii"] / self.p_inh - 1), .01)
        self.assertLess(abs(self.conn["p_ie"] / self.p_inh - 1), .01)

    def test_diag_zero(self):
        self.assertFalse(self.net.adjacency_matrix.diagonal().any())


class TestExtendedMiller(unittest.TestCase):
    def make_net(self, p_exc, p_inh, nr_neurons=2000):
        self.p_exc = p_exc
        self.p_inh = p_inh
        self.simple_miller = API(nr_neurons=2000, inh_ratio=.2, p_exc=self.p_exc,
                                 p_inh=self.p_inh, n_pow=4, feature_space_dimension=30)
        self.net = self.simple_miller()
        self.conn = ConnectivityEstimator(network=self.net)()

    def test_small_inh_connectivity(self):
        self.make_net(.2, .2)
        self.assertLess(abs(self.conn["p_ee"] / self.p_exc - 1), .01)
        self.assertLess(abs(self.conn["p_ei"] / self.p_exc - 1), .01)
        self.assertLess(abs(self.conn["p_ii"] / self.p_inh - 1), .015)
        self.assertLess(abs(self.conn["p_ie"] / self.p_inh - 1), .01)

    def test_small_population_small_inh_connectivity(self):
        self.make_net(.2, .2, 1000)
        self.assertLess(abs(self.conn["p_ee"] / self.p_exc - 1), .01)
        self.assertLess(abs(self.conn["p_ei"] / self.p_exc - 1), .01)
        self.assertLess(abs(self.conn["p_ii"] / self.p_inh - 1), .015)
        self.assertLess(abs(self.conn["p_ie"] / self.p_inh - 1), .01)

    def test_large_population_default_connectivity(self):
        self.make_net(.2, .5, 2000)
        self.assertLess(abs(self.conn["p_ee"] / self.p_exc - 1), .01)
        self.assertLess(abs(self.conn["p_ei"] / self.p_exc - 1), .01)
        self.assertLess(abs(self.conn["p_ii"] / self.p_inh - 1), .01)
        self.assertLess(abs(self.conn["p_ie"] / self.p_inh - 1), .01)

    def test_small_population_big_connectivity(self):
        self.make_net(.5, .5, 2000)
        self.assertLess(abs(self.conn["p_ee"] / self.p_exc - 1), .01)
        self.assertLess(abs(self.conn["p_ei"] / self.p_exc - 1), .01)
        self.assertLess(abs(self.conn["p_ii"] / self.p_inh - 1), .01)
        self.assertLess(abs(self.conn["p_ie"] / self.p_inh - 1), .01)



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