import unittest

from connectome.model import ER


class TestER(unittest.TestCase):
    def setUp(self):
        self.pe, self.pi = .3, .6
        self.ne = int(1500 / 3)
        self.ni = 1500 - self.ne
        self.er = ER(nr_neurons=1500, inh_ratio=1 / 3, p_exc=self.pe, p_inh=self.pi)

    def test_er(self):
        adj = self.er().adjacency_matrix
        ne, ni = self.ne, self.ni
        sum_e = adj[:,:ne].sum()
        sum_i = adj[:,-ni:].sum()
        exp_e = ne * (ne+ni) * self.pe
        exp_i = - ni * (ne+ni) * self.pi
        self.assertLess(abs(sum_e-exp_e)/exp_e, 0.01)
        self.assertLess(abs(sum_i-exp_i)/exp_i, 0.01)

    def test_diag_zero(self):
        adj = self.er().adjacency_matrix
        self.assertFalse(adj.diagonal().any())


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