import os
import os.path as path
import unittest
from multiprocessing import Pool
from tempfile import gettempdir

from connectome.builder import Sampler, InMemoryStorer, DistributedFileStorer


class MockModel:
    def __init__(self, value, name):
        self.value = value
        self.name = name

    def __call__(self):
        return self.value


class TestSampling(unittest.TestCase):
    def test_sampling_serial_map(self):
        persistence = InMemoryStorer()
        map_engine = map
        sampler = Sampler(map_engine, persistence)
        sampler.add_model(MockModel(1, "one"))
        sampler.add_model(MockModel(2, "two"))
        sampler.sample(nr_samples_per_model=7)
        results = sampler.results()
        self.assertEqual({"one": [1]*7,
                          "two": [2]*7}, results)


class TestMultiprocessingSampling(unittest.TestCase):
    def setUp(self):
        tmp_db = path.join(gettempdir(), "test.db")
        self.tmp_db = tmp_db
        self._delete_tmp()

    def tearDown(self):
        self._delete_tmp()

    def _delete_tmp(self):
        try:
            os.remove(self.tmp_db)
        except FileNotFoundError:
            pass

    def test_sampling_multiprocessing_map(self):
        persistence = DistributedFileStorer(self.tmp_db, timeout=5)
        with Pool(3) as p:
            map_engine = p.map
            sampler = Sampler(map_engine, persistence)
            sampler.add_model(MockModel(1, "one"))
            sampler.add_model(MockModel(2, "two"))
            sampler.sample(nr_samples_per_model=7)
        results = sampler.results()
        self.assertEqual(14, len(results))

    def test_sampling_multiprocessing_map_dict(self):
        persistence = DistributedFileStorer(self.tmp_db, timeout=5)
        with Pool(3) as p:
            map_engine = p.map
            sampler = Sampler(map_engine, persistence)
            sampler.add_model(MockModel({"one_res": 1}, "one"))
            sampler.add_model(MockModel({"two_res": 2}, "two"))
            sampler.sample(nr_samples_per_model=7)
        results = sampler.results()
        self.assertEqual(14, len(results))


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