import unittest
from math import sqrt

from connectome.analysis import ConnectivityEstimator, ReciprocityEstimator
from connectome.model import LL


class TestLL(unittest.TestCase):
    def setUp(self):
        self.p3 = {
        'nr_neurons': 2000,
        'inh_ratio': .25,
        'p_exc': .3,
        'p_inh': .6,
        'reciprocity_exc': .3,
        'nr_exc_subpopulations': 3}

        self.p2 = self.p3.copy()
        self.p2['nr_exc_subpopulations'] = 2

    def test_p_exc_3_layer(self):
        p = self.p3
        net = LL(**p)()
        ana = ConnectivityEstimator(network=net)()
        self.assertLess(abs(ana["p_ee"] - p["p_exc"]),  1 / sqrt(p["nr_neurons"]))
        self.assertLess(abs(ana["p_ee"] / p["p_exc"] - 1),  .01)
        self.assertLess(abs(ana["p_ei"] - p["p_exc"]),  1 / sqrt(p["nr_neurons"]))
        self.assertLess(abs(ana["p_ei"] / p["p_exc"] - 1),  .01)

    def test_p_inh_3_layer(self):
        p = self.p3
        net = LL(**p)()
        ana = ConnectivityEstimator(network=net)()
        self.assertLess(abs(ana["p_ii"] - p["p_inh"]),  1 / sqrt(p["nr_neurons"]))
        self.assertLess(abs(ana["p_ii"] / p["p_inh"] - 1),  .01)
        self.assertLess(abs(ana["p_ie"] - p["p_inh"]),  1 / sqrt(p["nr_neurons"]))
        self.assertLess(abs(ana["p_ie"] / p["p_inh"] - 1),  .01)

    def test_p_exc_2_layer(self):
        p = self.p2
        net = LL(**p)()
        ana = ConnectivityEstimator(network=net)()
        self.assertLess(abs(ana["p_ee"] - p["p_exc"]),  1 / sqrt(p["nr_neurons"]))
        self.assertLess(abs(ana["p_ee"] / p["p_exc"] - 1),  .01)
        self.assertLess(abs(ana["p_ei"] - p["p_exc"]),  1 / sqrt(p["nr_neurons"]))
        self.assertLess(abs(ana["p_ei"] / p["p_exc"] - 1),  .01)

    def test_reciprocity_exc_2_layer(self):
        p = self.p2
        net = LL(**p)()
        ana = ReciprocityEstimator(network=net)()
        r_e = ana["reciprocity_ee"]
        self.assertLess(abs(r_e - p["reciprocity_exc"]), 1 / sqrt(p["nr_neurons"]))
        self.assertLess(abs(r_e / p["reciprocity_exc"] - 1), .01)

    def test_reciprocity_exc_3_layer(self):
        p = self.p2
        net = LL(**p)()
        ana = ReciprocityEstimator(network=net)()
        r_e = ana["reciprocity_ee"]
        self.assertLess(abs(r_e - p["reciprocity_exc"]), 1 / sqrt(p["nr_neurons"]))
        self.assertLess(abs(r_e / p["reciprocity_exc"] - 1), .01)

    def test_diag_zero(self):
        p = self.p3
        net = LL(**p)()
        self.assertFalse(net.adjacency_matrix.diagonal().any())

    def test_nr_blocks_2(self):
        net = LL(**self.p2)()
        self.assertEqual(3, len(net.blocks))
        self.assertEqual(2, net.blocks.nr_exc)

    def test_nr_blocks_3(self):
        net = LL(**self.p3)()
        self.assertEqual(4, len(net.blocks))
        self.assertEqual(3, net.blocks.nr_exc)
