import unittest
from connectome.model import ERFEVERInitialConnectivity, EXPFEVERInitialConnectivity, Network, ERFEVER, EXPFEVER

import scipy as sp

from connectome.analysis import ConnectivityEstimator, ReciprocityEstimator, RelativeReciprocityEstimator
from connectome.model.fever import FEVERGivenInitialGraph
from util.testing import run_slow_test


class TestErFeverInitialConnectivity(unittest.TestCase):
    def test_connectivity_and_reciprocity(self):
        p_exc = 0.2
        p_inh = 0.4
        nr_exc = 300
        p = {"nr_neurons": 600, "inh_ratio": .5, "p_exc": p_exc, "p_inh": p_inh, "feature_space_dimension": 30,
             "p_initial_graph_exc": .9*p_exc, "p_initial_graph_inh": .97*p_inh}
        net = ERFEVERInitialConnectivity()(p)
        W = net.adjacency_matrix
        n = ConnectivityEstimator()(network=net)
        r = ReciprocityEstimator()(network=net)
        rr = RelativeReciprocityEstimator()(network=net)
        p_e, p_i = n['p_ee'], n['p_ii']
        self.assertLess(abs(p_e / p_exc -1), .1)
        self.assertLess(abs(p_i / p_inh -1), .1)
        self.assertFalse(W.diagonal().any())
        self.assertTrue((W[:,:nr_exc] >= 0).all())
        self.assertTrue((W[:,nr_exc:] <= 0).all())
        self.assertGreaterEqual(r['reciprocity_ee'] / p_e, 1.1)
        self.assertLessEqual(r['reciprocity_ee'] / p_e, 2)
        self.assertLess(rr['relative_reciprocity_ei'], 1.05)
        self.assertLess(rr['relative_reciprocity_ie'], 1.05)

    @unittest.skipUnless(run_slow_test(), "Slow test")
    def test_very_small_connectivity_zero_initial_connectivity(self):
        p_exc = 0.02
        p_inh = 0.05
        nr_exc = 300
        p = {"nr_neurons": 600, "inh_ratio": .5, "p_exc": p_exc, "p_inh": p_inh, "feature_space_dimension": 30,
             "p_initial_graph_exc": 0, "p_initial_graph_inh": 0}
        net = ERFEVERInitialConnectivity()(p)
        W = net.adjacency_matrix
        n = ConnectivityEstimator()(network=net)
        r = ReciprocityEstimator()(network=net)
        rr = RelativeReciprocityEstimator()(network=net)
        p_e, p_i = n['p_ee'], n['p_ii']
        self.assertLess(abs(p_e / p_exc -1), .1)
        self.assertLess(abs(p_i / p_inh -1), .1)
        self.assertFalse(W.diagonal().any())
        self.assertTrue((W[:,:nr_exc] >= 0).all())
        self.assertTrue((W[:,nr_exc:] <= 0).all())
        self.assertGreaterEqual(r['reciprocity_ee'] / p_e, 1.2)
        self.assertLess(rr['relative_reciprocity_ei'], 1.05)
        self.assertLess(rr['relative_reciprocity_ie'], 1.05)


class TestExpFeverInitialConnectivity(unittest.TestCase):
    def test_connectivity_and_reciprocity(self):
        p_exc = 0.2
        p_inh = 0.4
        nr_exc = 300
        p = {"nr_neurons": 600, "inh_ratio": .5, "p_exc": p_exc, "p_inh": p_inh, "feature_space_dimension": 30,
             "p_initial_graph_exc": .9*p_exc, "p_initial_graph_inh": .97*p_inh}
        net = EXPFEVERInitialConnectivity()(p)
        W = net.adjacency_matrix
        n = ConnectivityEstimator()(network=net)
        r = ReciprocityEstimator()(network=net)
        rr = RelativeReciprocityEstimator()(network=net)
        p_e, p_i = n['p_ee'], n['p_ii']
        self.assertLess(abs(p_e / p_exc -1), .1)
        self.assertLess(abs(p_i / p_inh -1), .1)
        self.assertFalse(W.diagonal().any())
        self.assertTrue((W[:,:nr_exc] >= 0).all())
        self.assertTrue((W[:,nr_exc:] <= 0).all())
        self.assertGreaterEqual(r['reciprocity_ee'] / p_e, 1.1)
        self.assertLessEqual(r['reciprocity_ee'] / p_e, 2)
        self.assertGreaterEqual(rr['relative_reciprocity_ei'], 1.1)
        self.assertGreaterEqual(rr['relative_reciprocity_ie'], 1.1)


class TestAlmostZeroInitialNetowrk(unittest.TestCase):
    def test_no_zero_division_error(self):
        initial = sp.zeros((10,10))
        initial[2,3] = .2
        FEVERGivenInitialGraph(initial_network=Network(initial),
                               feature_vectors=sp.rand(*initial.shape),
                               p_exc=.05,
                               p_inh=.05)()

    def test_small_network_no_error(self):
        initial = sp.zeros((5,5))
        initial[1,0] = .2
        FEVERGivenInitialGraph(initial_network=Network(initial),
                               feature_vectors=sp.rand(*initial.shape),
                               p_exc=.05,
                               p_inh=.05)()


class TestERFEver(unittest.TestCase):
    def test_connectivity_and_reciprocity(self):
        p_exc = 0.2
        p_inh = 0.4
        p = {"nr_neurons": 600, "inh_ratio": .5, "p_exc": p_exc, "p_inh": p_inh, "feature_space_dimension": 30,
             "feverization_ratio": .7}
        erfever = ERFEVER(**p)
        net = erfever()
        W = net.adjacency_matrix
        n = ConnectivityEstimator()(network=net)
        r = ReciprocityEstimator()(network=net)
        rr = RelativeReciprocityEstimator()(network=net)
        p_e, p_i = n['p_ee'], n['p_ii']
        self.assertLess(abs(p_e / p_exc -1), .05)
        self.assertLess(abs(p_i / p_inh -1), .05)
        self.assertFalse(W.diagonal().any())
        self.assertTrue((W[:,:net.nr_exc] >= 0).all())
        self.assertTrue((W[:,net.nr_exc:] <= 0).all())
        self.assertGreaterEqual(r['reciprocity_ee'] / p_e, 1.1)
        self.assertLessEqual(r['reciprocity_ee'] / p_e, 2)
        self.assertLess(rr['relative_reciprocity_ei'], 1.05)
        self.assertLess(rr['relative_reciprocity_ie'], 1.05)

    def test_get_lambdas(self):
        """
        This test is some how stuped.
        It tries to ensure that the get_lamdbas method works.
        The assertGreater is not very meaningful.
        """
        erfever = ERFEVER(nr_neurons=400, inh_ratio=.5, p_exc=.2, p_inh=.5, feature_space_dimension=100,
                         feverization_ratio=.7)
        erfever()
        lambda_exc, lambda_inh = erfever.get_lambdas()
        self.assertGreater(lambda_exc, 0)
        self.assertGreater(lambda_inh, 0)


class TestEXPFEver(unittest.TestCase):
    def test_connectivity_and_reciprocity(self):
        p_exc = 0.2
        p_inh = 0.4
        nr_exc = 300
        p = {"nr_neurons": 600, "inh_ratio": .5, "p_exc": p_exc, "p_inh": p_inh, "feature_space_dimension": 30,
             "feverization_ratio": .7}
        net = EXPFEVER()(p)
        self.assertEqual(300, net.nr_exc)
        self.assertEqual(300, net.nr_inh)
        W = net.adjacency_matrix
        n = ConnectivityEstimator()(network=net)
        r = ReciprocityEstimator()(network=net)
        rr = RelativeReciprocityEstimator()(network=net)
        p_e, p_i = n['p_ee'], n['p_ii']
        self.assertLess(abs(p_e / p_exc -1), .1)
        self.assertLess(abs(p_i / p_inh -1), .1)
        self.assertFalse(W.diagonal().any())
        self.assertTrue((W[:,:nr_exc] >= 0).all())
        self.assertTrue((W[:,nr_exc:] <= 0).all())
        self.assertGreaterEqual(r['reciprocity_ee'] / p_e, 1.1)
        self.assertLessEqual(r['reciprocity_ee'] / p_e, 2)
        self.assertGreaterEqual(rr['relative_reciprocity_ei'], 1.1)
        self.assertGreaterEqual(rr['relative_reciprocity_ie'], 1.1)

    def test_synapse_sign(self):
        p_exc = 0.10155
        p_inh = 0.482
        nr_neurons = 50 + 20
        inh_ratio = 20 / nr_neurons
        p = {"nr_neurons": nr_neurons, "inh_ratio": inh_ratio, "p_exc": p_exc, "p_inh": p_inh,
             "feature_space_dimension": 4,
             "feverization_ratio": .38}
        net = EXPFEVER()(p)
        W = net.adjacency_matrix

    def test_get_lambdas(self):
        """
        This test is some how stuped.
        It tries to ensure that the get_lamdbas method works.
        The assertGreater is not very meaningful.
        """
        expfever = EXPFEVER(nr_neurons=400, inh_ratio=.5, p_exc=.2, p_inh=.5, feature_space_dimension=100,
                         feverization_ratio=.7)
        expfever()
        lambda_exc, lambda_inh = expfever.get_lambdas()
        self.assertGreater(lambda_exc, 0)
        self.assertGreater(lambda_inh, 0)


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