import unittest

import scipy as sp

from connectome.analysis import ConnectivityEstimator, RelativeReciprocityEstimator
from connectome.model import EXP
from connectome.noise import OutDegreePreservingNoise


class TestExpDecaying(unittest.TestCase):
    def setUp(self):
        self.p_exc = .2
        self.p_inh = .52
        decay = sp.rand() * .5 + .5
        self.exp = EXP(nr_neurons=2000, inh_ratio=.2, p_exc=.2, p_inh=self.p_inh, decay=decay)
        self.net = self.exp()
        self.connectivity = ConnectivityEstimator(network=self.net)()
        self.relative_reciprocity = RelativeReciprocityEstimator(network=self.net)()

    def test_exc_connectivity(self):
        p = self.connectivity["p_ee"]
        self.assertLess(abs(p/self.p_exc - 1), .05)

    def test_inh_connectivity(self):
        pi = self.connectivity["p_ii"]
        self.assertLess(abs(pi/self.p_inh - 1), .05)

    def test_diag_zero(self):
        self.assertFalse(sp.diag(self.net.adjacency_matrix).any())

    def test_reciprocity_no_noise(self):
        self.assertGreater(self.relative_reciprocity["relative_reciprocity_ee"], 1.1)

    def test_reshuffle_noise(self):
        noisy_net = OutDegreePreservingNoise(network=self.net, fraction_draws_over_nr_synapses=1)()
        rec = RelativeReciprocityEstimator(network=noisy_net)()["relative_reciprocity_ee"]
        self.assertLess(abs(rec-1), .1)


class TestP0Variation(unittest.TestCase):
    def test_p_0_variation(self):
        exp_low_rec = EXP(nr_neurons=2000, inh_ratio=.2, p_exc=.2, p_inh=.6, decay=.3)
        exp_high_rec = EXP(nr_neurons=2000, inh_ratio=.2, p_exc=.2, p_inh=.6, decay=1)
        net_low_rec = exp_low_rec()
        net_high_rec = exp_high_rec()
        rec_low = RelativeReciprocityEstimator(network=net_low_rec)()["relative_reciprocity_ee"]
        rec_high = RelativeReciprocityEstimator(network=net_high_rec)()["relative_reciprocity_ee"]
        self.assertGreater(rec_high, rec_low)

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