"""
The tests here are mostly integration test.
Trying to make sure I don't break the execution.
But don't thake the criteria seriously. They are just mocked.
"""

import unittest

from brian2 import ms

from connectome.function import TestSuite, Range, TextureTask, GreaterThan, LessThan, UnsynchronizedActivityTask, MemoryTask, TuningTask, PropagationTask
from connectome.model import API, SORN, ERFEVER, SYN, ER, LL, EXP


class TestValidation(unittest.TestCase):
    def test_pipeline_integration(self):
        test = TestSuite(lambda x: {"property1" : x, "property2": x * 2})
        test.add_criterion("property1", Range(1, 2))
        test.add_criterion("property2", Range(2, 3))
        test_result = test.test_model(1)
        self.assertTrue(test_result.passed)


class TestTextureTask(unittest.TestCase):
    nr_epochs = 2

    def suite_factory(self):
        return TestSuite(TextureTask(nb_epoch=self.nr_epochs, nr_samples=10, nr_classes=2, length=10, verbose=2))

    def test__ll(self):
        test_suite = self.suite_factory()
        test_suite.add_criterion("acc", GreaterThan(.6))

        ll = LL(nr_neurons=150, inh_ratio=.2, p_exc=.2, p_inh=.5, nr_exc_subpopulations=2, reciprocity_exc=.2)
        test_result = test_suite.test_model(ll)
        self.assertTrue(test_result.passed or not test_result.passed)  # just check if no errors

    def test_er(self):
        test_suite = self.suite_factory()
        test_suite.add_criterion("acc", GreaterThan(.6))

        er = ER(nr_neurons=50, inh_ratio=.5, p_exc=.2, p_inh=.5)
        test_result = test_suite.test_model(er)
        self.assertTrue(test_result.passed or not test_result.passed)  # just check if no errors

    def test_exp(self):
        test_suite = self.suite_factory()
        test_suite.add_criterion("acc", GreaterThan(.6))

        er = EXP(nr_neurons=50, inh_ratio=.5, p_exc=.2, p_inh=.5, decay=1)
        test_result = test_suite.test_model(er)
        self.assertTrue(test_result.passed or not test_result.passed)  # just check if no errors

    def test_syn(self):
        test_suite = self.suite_factory()
        test_suite.add_criterion("acc", GreaterThan(.6))

        syn = SYN(nr_neurons=50, inh_ratio=.5, p_exc=.2, p_inh=.5, pool_size=10)
        test_result = test_suite.test_model(syn)
        self.assertTrue(test_result.passed or not test_result.passed)  # just check if no errors

    def test_non_checked_criterion_is_returned(self):
        test_suite = self.suite_factory()
        test_suite.add_criterion("acc", GreaterThan(.6))

        er = ER(nr_neurons=50, inh_ratio=.5, p_exc=.2, p_inh=.5)
        test_result = test_suite.test_model(er)
        self.assertTrue(isinstance(test_result["id"].score, int))
        # also check that the pass works
        self.assertTrue(test_result.passed or not test_result.passed)

    def test_non_checked_criterion_passed_still_works(self):
        test_suite = self.suite_factory()
        test_suite.add_criterion("acc", GreaterThan(.6))

        er = ER(nr_neurons=50, inh_ratio=.5, p_exc=.2, p_inh=.5)
        test_result = test_suite.test_model(er)
        self.assertTrue(test_result.passed or not test_result.passed)


class TestMemoryTask(unittest.TestCase):
    def test_fever(self):
        test_suite = TestSuite(MemoryTask())
        test_suite.add_criterion("cosine_similarity", GreaterThan(.1))

        erfever = ERFEVER(nr_neurons=20, inh_ratio=.4, p_exc=.2, p_inh=.2, feature_space_dimension=20,
                          feverization_ratio=.6)
        res = test_suite.test_model(erfever)
        self.assertTrue(res.passed or not res.passed)

    def test_api(self):
        test_suite = TestSuite(MemoryTask())
        test_suite.add_criterion("cosine_similarity", GreaterThan(-2))

        api = API(nr_neurons=20, inh_ratio=.4, p_exc=.2, p_inh=.2, feature_space_dimension=10, n_pow=4)
        res = test_suite.test_model(api)
        self.assertEqual(4, res["n_pow"].score)
        self.assertTrue(res.passed or not res.passed)

    def test_api_default_output(self):
        test_suite = TestSuite(MemoryTask())

        api = API(nr_neurons=20, inh_ratio=.4, p_exc=.2, p_inh=.2, feature_space_dimension=10, n_pow=4)
        res = test_suite.test_model(api)
        self.assertEqual(4, res["n_pow"].score)
        self.assertEqual(10, res["feature_space_dimension"].score)
        self.assertEqual(20, res["nr_neurons"].score)
        self.assertEqual("API", res["model_name"].score)


class TestTuningTask(unittest.TestCase):
    def test_api(self):
        test_suite = TestSuite(TuningTask())
        test_suite.add_criterion("pool_activity_min", LessThan(0))
        test_suite.add_criterion("pool_activity_max", GreaterThan(0))

        api = API(nr_neurons=20, inh_ratio=.2, p_exc=.3, p_inh=.5, feature_space_dimension=20, n_pow=5)
        res = test_suite.test_model(api)
        self.assertTrue(res["pool_activity_max"].score is not None)
        self.assertTrue(res["pool_activity_min"].score is not None)
        self.assertTrue(res.passed)

    def test_fever(self):
        test_suite = TestSuite(TuningTask())
        test_suite.add_criterion("pool_activity_min", LessThan(0))
        test_suite.add_criterion("pool_activity_max", GreaterThan(0))

        erfever = ERFEVER(nr_neurons=20, inh_ratio=.4, p_exc=.2, p_inh=.2, feature_space_dimension=20,
                          feverization_ratio=.6)
        res = test_suite.test_model(erfever)
        self.assertTrue(res["pool_activity_max"].score is not None)
        self.assertTrue(res["pool_activity_min"].score is not None)
        self.assertTrue(res.passed)


class TestUnsynchronizedActivityTask(unittest.TestCase):
    def test_sorn(self):
        sorn_test = TestSuite(UnsynchronizedActivityTask())
        sorn_test.add_criterion("correlation", Range(-1.1, 1.1))

        sorn = SORN(nr_neurons=200, inh_ratio=.1, p_exc=.2, p_inh=.4, eta_stdp=.001, eta_intrinsic=.001)
        res = sorn_test.test_model(sorn)
        self.assertTrue(res.passed)


class TestropagationTask(unittest.TestCase):
    def test_syn(self):
        syn_test = TestSuite(PropagationTask(100 * ms))
        syn_test.add_criterion("fraction_pools_activated", Range(0, 1))

        syn = SYN(nr_neurons=2000, inh_ratio=.1, p_exc=.15, p_inh=.5, pool_size=100)
        res = syn_test.test_model(syn)
        self.assertTrue(res.passed)


class TestBatchProcessing(unittest.TestCase):
    def test_batch(self):
        def doubler(x):
            return {"res2": 2*x, "res": x}
        test_suite = TestSuite(doubler)
        test_suite.add_criterion("res", Range(0, 1))
        test_suite.add_criterion("res2", GreaterThan(.8))
        test_suite.enqueue_model(.3)
        test_suite.enqueue_model(.5)
        test_suite.enqueue_model(.7)
        test_suite.enqueue_model(1.7)
        results = test_suite.execute_enqueued_models()
        passed = [bool(res) for res in results]
        self.assertEqual([False, True, True, False], passed)


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