import unittest

import scipy as sp
from connectome.model.dynamicsorn.fast_cython_sorn import (make_sorn)

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


class TestUpperVsLowerTriangular(unittest.TestCase):
    def test_upper_and_lower_connectivity_fraction_equal(self):
        nr_exc = 600
        nr_inh = 200
        nr_time_steps = 600
        p_exc = .2
        p_inh = .2
        eta_equal = .2
        eta_plus = .2
        eta_minus = -.2
        adj = make_sorn(nr_time_steps, nr_exc, nr_inh, p_exc, p_inh, eta_equal, eta_plus, eta_minus)
        upper = sp.triu(adj[:nr_exc,:nr_exc], 1)
        lower = sp.tril(adj[:nr_exc,:nr_exc], -1)
        self.assertLess(abs(upper.sum() - lower.sum()) / (upper.sum() + lower.sum()), .07)


class TestSORNPLausibleValues(unittest.TestCase):
    def test_plausible_sorn(self):
        class Recorder:
            def __init__(self, delta_t=1):
                self.r_ee = []
                self.p_ee = []
                self.activity = []
                self.delta_t = delta_t
                self.t = []
                self.last_step = -1
                self.firing_threshold_max = []
                self.firing_threshold_min = []
                self.firing_threshold_mean = []
                self.firing_threshold_median = []
                self.firing_threshold_std = []

            def record(self, t, name, value):
                if not t % self.delta_t:
                    if self.last_step != t:
                        self.t.append(t)
                        self.last_step = t
                    if name == "adjacency_matrix":
                        net = Network(value)
                        self.r_ee.append(ReciprocityEstimator(network=net)()["reciprocity_ee"])
                        self.p_ee.append(ConnectivityEstimator(network=net)()["p_ee"])
                    if name == "current_active_neurons":
                        self.activity.append(value.sum() / value.size)
                    if name == "firing_threshold":
                        self.firing_threshold_max.append(value.max())
                        self.firing_threshold_min.append(value.min())
                        self.firing_threshold_mean.append(value.mean())
                        self.firing_threshold_median.append(sp.median(value))
                        self.firing_threshold_std.append(sp.sqrt(sp.var(value)))

        recorder = Recorder(delta_t=100)
        p_exc = .2
        adj = make_sorn(p_exc=p_exc, nr_time_steps=2000, nr_exc=1200, nr_inh=1200//5, history_recorder=recorder)
        # this test has been added after exploration of SORN to make decrease the chance
        # of erroneous code changes breaking the model
        self.assertLess(max(recorder.activity), .8)
        self.assertLess(min(recorder.r_ee), p_exc * .8)
        self.assertFalse(adj.diagonal().any())


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