import types
import unittest
from itertools import chain, product
from multiprocessing import cpu_count
from unittest.mock import patch

import numpy as np

from mpire.utils import apply_numpy_chunking, chunk_tasks, format_seconds, get_n_chunks, make_single_arguments, TimeIt


class ChunkTasksTest(unittest.TestCase):

    def test_no_chunk_size_no_n_splits_provided(self):
        """
        Test that a ValueError is raised when no chunk_size and n_splits are provided
        """
        with self.assertRaises(ValueError):
            next(chunk_tasks([]))

    def test_generator_without_iterable_len(self):
        """
        Test that a ValueError is raised when a generator is provided without iterable_len
        """
        with self.assertRaises(ValueError):
            next(chunk_tasks(iter([]), n_splits=1))

    def test_chunk_size_has_priority_over_n_splits(self):
        """
        Test that chunk_size is prioritized over n_splits
        """
        chunks = list(chunk_tasks(range(4), chunk_size=4, n_splits=4))
        self.assertEqual(len(chunks), 1)
        self.assertEqual(len(chunks[0]), 4)
        self.assertEqual(list(range(4)), list(chain.from_iterable(chunks)))

    def test_empty_input(self):
        """
        Test that the chunker is an empty generator for an empty input iterable
        """
        with self.subTest('list input'):
            chunks = list(chunk_tasks([], n_splits=5))
            self.assertEqual(len(chunks), 0)

        with self.subTest('generator/iterator input'):
            chunks = list(chunk_tasks(iter([]), iterable_len=0, n_splits=5))
            self.assertEqual(len(chunks), 0)

    def test_iterable_len_doesnt_match_input_size(self):
        """
        Test for cases where iterable_len does and does not match the number of arguments (it should work fine)
        """
        num_args = 10
        for iter_len in [5, 10, 20]:
            expected_args_sum = min(iter_len, num_args)

            # Test for normal list (range is considered a normal list as it implements __len__ and such)
            with self.subTest(iter_len=iter_len, input='list'):
                chunks = list(chunk_tasks(range(num_args), iterable_len=iter_len, n_splits=1))
                total_args = sum(map(len, chunks))
                self.assertEqual(total_args, expected_args_sum)
                self.assertEqual(list(range(expected_args_sum)), list(chain.from_iterable(chunks)))

            # Test for an actual generator (range does not really behave like one)
            with self.subTest(iter_len=iter_len, input='generator/iterator'):
                chunks = list(chunk_tasks(iter(range(num_args)), iterable_len=iter_len, n_splits=1))
                total_args = sum(map(len, chunks))
                self.assertEqual(total_args, expected_args_sum)
                self.assertEqual(list(range(expected_args_sum)), list(chain.from_iterable(chunks)))

    def test_n_splits(self):
        """
        Test different values of n_splits: len(args) {<, ==, >} n_splits
        """
        n_splits = 5
        for num_args in [n_splits - 1, n_splits, n_splits + 1]:
            expected_n_chunks = min(n_splits, num_args)

            # Test for normal list (range is considered a normal list as it implements __len__ and such)
            with self.subTest(num_args=num_args, input='list'):
                chunks = list(chunk_tasks(range(num_args), n_splits=n_splits))
                self.assertEqual(len(chunks), expected_n_chunks)
                self.assertEqual(list(range(num_args)), list(chain.from_iterable(chunks)))

            # Test for an actual generator (range does not really behave like one)
            with self.subTest(num_args=num_args, input='generator/iterator'):
                chunks = list(chunk_tasks(iter(range(num_args)), iterable_len=num_args, n_splits=n_splits))
                self.assertEqual(len(chunks), expected_n_chunks)
                self.assertEqual(list(range(num_args)), list(chain.from_iterable(chunks)))

    def test_chunk_size(self):
        """
        Test that chunks are of the right size if chunk_size is provided
        """
        chunk_size = 3
        for num_args in [chunk_size - 1, chunk_size, chunk_size + 1]:

            # Test for normal list (range is considered a normal list as it implements __len__ and such)
            with self.subTest(num_args=num_args, input='list'):
                chunks = list(chunk_tasks(range(num_args), chunk_size=chunk_size))
                for chunk in chunks[:-1]:
                    self.assertEqual(len(chunk), chunk_size)
                self.assertLessEqual(len(chunks[-1]), chunk_size)
                self.assertEqual(list(range(num_args)), list(chain.from_iterable(chunks)))

            # Test for an actual generator (range does not really behave like one)
            with self.subTest(num_args=num_args, input='generator/iterator'):
                chunks = list(chunk_tasks(iter(range(num_args)), chunk_size=chunk_size))
                for chunk in chunks[:-1]:
                    self.assertEqual(len(chunk), chunk_size)
                self.assertLessEqual(len(chunks[-1]), chunk_size)
                self.assertEqual(list(range(num_args)), list(chain.from_iterable(chunks)))


class ApplyNumpyChunkingTest(unittest.TestCase):

    """
    This function simply calls other, already tested, functions in succession. We do test the individual parameter
    influence, but interactions between them are skipped
    """

    def setUp(self):
        self.test_data_numpy = np.random.rand(100, 2)

    def test_iterable_len(self):
        """
        Test that iterable_len is adhered to. When iterable_len < len(input) it should reduce the input size. If higher
        or None it should take the entire input
        """
        for iterable_len, expected_size in [(5, 5), (150, 100), (None, 100)]:
            with self.subTest(iterable_len=iterable_len):
                iterable_of_args, iterable_len_, chunk_size, n_splits = apply_numpy_chunking(
                    self.test_data_numpy, iterable_len=iterable_len, n_splits=1
                )

                # Materialize generator and test contents
                iterable_of_args = list(iterable_of_args)
                self.assertEqual(len(iterable_of_args), 1)
                self.assertIsInstance(iterable_of_args[0][0], np.ndarray)
                np.testing.assert_array_equal(iterable_of_args[0][0], self.test_data_numpy[:expected_size])

                # Test other output
                self.assertEqual(iterable_len_, 1)
                self.assertEqual(chunk_size, 1)
                self.assertIsNone(n_splits)

    def test_chunk_size(self):
        """
        Test that chunk_size works as expected. Note that chunk_size trumps n_splits
        """
        for chunk_size, expected_n_chunks in [(1, 100), (3, 34), (200, 1), (None, 1)]:
            with self.subTest(chunk_size=chunk_size):
                iterable_of_args, iterable_len, chunk_size_, n_splits = apply_numpy_chunking(
                    self.test_data_numpy, chunk_size=chunk_size, n_splits=1
                )

                # Materialize generator and test contents. The chunks should be of size chunk_size (expect for the last
                # chunk which can be smaller)
                iterable_of_args = list(iterable_of_args)
                self.assertEqual(len(iterable_of_args), expected_n_chunks)
                chunk_size = chunk_size or 100
                for chunk_idx, chunk in enumerate(iterable_of_args):
                    self.assertIsInstance(chunk[0], np.ndarray)
                    np.testing.assert_array_equal(chunk[0], self.test_data_numpy[chunk_idx * chunk_size:
                                                                                 (chunk_idx + 1) * chunk_size])

                # Test other output
                self.assertEqual(iterable_len, expected_n_chunks)
                self.assertEqual(chunk_size_, 1)
                self.assertIsNone(n_splits)

    def test_n_splits(self):
        """
        Test that n_splits works as expected.
        """
        for n_splits, expected_n_chunks in [(1, 1), (3, 3), (150, 100)]:
            with self.subTest(n_splits=n_splits):
                iterable_of_args, iterable_len, chunk_size, n_splits_ = apply_numpy_chunking(
                    self.test_data_numpy, n_splits=n_splits
                )

                # Materialize generator and test contents. We simply test if every row of the original input occurs in
                # the chunks
                iterable_of_args = list(iterable_of_args)
                self.assertEqual(len(iterable_of_args), expected_n_chunks)
                offset = 0
                for chunk in iterable_of_args:
                    self.assertIsInstance(chunk[0], np.ndarray)
                    np.testing.assert_array_equal(chunk[0], self.test_data_numpy[offset:offset + len(chunk[0])])
                    offset += len(chunk[0])
                self.assertEqual(offset, 100)

                # Test other output
                self.assertEqual(iterable_len, expected_n_chunks)
                self.assertEqual(chunk_size, 1)
                self.assertIsNone(n_splits_)

        # chunk_size and n_splits can't be both None
        with self.subTest(n_splits=None), self.assertRaises(ValueError):
            iterable_of_args, *_ = apply_numpy_chunking(self.test_data_numpy, n_splits=None)
            list(iterable_of_args)

    def test_n_jobs(self):
        """
        Test that n_jobs works as expected. When chunk_size and n_splits are both None, n_jobs * 4 is passed on as
        n_splits
        """
        for n_jobs, expected_n_chunks in [(1, 4), (3, 12), (40, 100), (150, 100)]:
            with self.subTest(n_jobs=n_jobs):
                iterable_of_args, iterable_len, chunk_size, n_splits_ = apply_numpy_chunking(
                    self.test_data_numpy, n_jobs=n_jobs
                )

                # Materialize generator and test contents. We simply test if every row of the original input occurs in
                # the chunks
                iterable_of_args = list(iterable_of_args)
                self.assertEqual(len(iterable_of_args), expected_n_chunks)
                offset = 0
                for chunk in iterable_of_args:
                    self.assertIsInstance(chunk[0], np.ndarray)
                    np.testing.assert_array_equal(chunk[0], self.test_data_numpy[offset:offset + len(chunk[0])])
                    offset += len(chunk[0])
                self.assertEqual(offset, 100)

                # Test other output
                self.assertEqual(iterable_len, expected_n_chunks)
                self.assertEqual(chunk_size, 1)
                self.assertIsNone(n_splits_)


class GetNChunksTest(unittest.TestCase):

    def setUp(self):
        self.test_data = [1, 2, 3, 5, 6, 9, 37, 42, 1337, 0, 3, 5, 0]
        self.test_data_numpy = np.random.rand(100, 2)

    def test_everything_none(self):
        """
        When everything is None we should use cpu_count * 4 as number of splits. We have to take the number of tasks
        into account
        """
        with self.subTest(input='list'):
            self.assertEqual(get_n_chunks(self.test_data, iterable_len=None, chunk_size=None, n_splits=None,
                                          n_jobs=None), min(13, cpu_count() * 4))
        with self.subTest(input='numpy'):
            self.assertEqual(get_n_chunks(self.test_data_numpy, iterable_len=None, chunk_size=None, n_splits=None,
                                          n_jobs=None), min(100, cpu_count() * 4))

    def test_smaller_iterable_len(self):
        """
        Test iterable_len, where iterable_len < len(input)
        """
        with self.subTest(input='list'):
            self.assertEqual(get_n_chunks(self.test_data, iterable_len=5, chunk_size=None, n_splits=None, n_jobs=None),
                             min(5, cpu_count() * 4))
        with self.subTest(input='numpy'):
            self.assertEqual(get_n_chunks(self.test_data_numpy, iterable_len=5, chunk_size=None, n_splits=None,
                                          n_jobs=None), min(5, cpu_count() * 4))
        with self.subTest(input='generator/iterator'):
            self.assertEqual(get_n_chunks(iter(self.test_data), iterable_len=5, chunk_size=None, n_splits=None,
                                          n_jobs=None), min(5, cpu_count() * 4))

    def test_larger_iterable_len(self):
        """
        Test iterable_len, where iterable_len > len(input). Should ignores iterable_len when actual number of tasks is
        less, except when we use the data_generator function, in which case we cannot determine the actual number of
        elements.
        """
        with self.subTest(input='list'):
            self.assertEqual(get_n_chunks(self.test_data, iterable_len=25, chunk_size=None, n_splits=None, n_jobs=None),
                             min(13, cpu_count() * 4))
        with self.subTest(input='numpy'):
            self.assertEqual(get_n_chunks(self.test_data_numpy, iterable_len=125, chunk_size=None, n_splits=None,
                                          n_jobs=None), min(100, cpu_count() * 4))
        with self.subTest(input='generator/iterator'):
            self.assertEqual(get_n_chunks(iter(self.test_data), iterable_len=25, chunk_size=None, n_splits=None,
                                          n_jobs=None), min(25, cpu_count() * 4))

    def test_chunk_size(self):
        """
        Test chunk_size
        """
        for chunk_size, expected_n_chunks in [(1, 13), (3, 5)]:
            with self.subTest(input='list', chunk_size=chunk_size, expected_n_chunks=expected_n_chunks):
                self.assertEqual(get_n_chunks(self.test_data, iterable_len=None, chunk_size=chunk_size, n_splits=None,
                                              n_jobs=None), expected_n_chunks)

        for chunk_size, expected_n_chunks in [(1, 100), (3, 34)]:
            with self.subTest(input='list', chunk_size=chunk_size, expected_n_chunks=expected_n_chunks):
                self.assertEqual(get_n_chunks(self.test_data_numpy, iterable_len=None, chunk_size=chunk_size,
                                              n_splits=None, n_jobs=None), expected_n_chunks)

    def test_n_splits(self):
        """
        Test n_splits. n_jobs shouldn't have any influence
        """
        for n_splits, n_jobs in product([1, 6], [None, 2, 8]):
            with self.subTest(input='list', n_splits=n_splits, n_jobs=n_jobs):
                self.assertEqual(get_n_chunks(self.test_data, iterable_len=None, chunk_size=None, n_splits=n_splits,
                                              n_jobs=n_jobs), n_splits)

            with self.subTest(input='numpy', n_splits=n_splits, n_jobs=n_jobs):
                self.assertEqual(get_n_chunks(self.test_data_numpy, iterable_len=None, chunk_size=None,
                                              n_splits=n_splits, n_jobs=n_jobs), n_splits)

    def test_n_jobs(self):
        """
        When everything is None except n_jobs we should use n_jobs * 4 as number of splits. Again, taking into account
        the number of tasks
        """
        for n_jobs in [1, 6]:
            with self.subTest(input='list', n_jobs=n_jobs):
                self.assertEqual(get_n_chunks(self.test_data, iterable_len=None, chunk_size=None, n_splits=None,
                                              n_jobs=n_jobs), min(4 * n_jobs, len(self.test_data)))

            with self.subTest(input='numpy', n_jobs=n_jobs):
                self.assertEqual(get_n_chunks(self.test_data_numpy, iterable_len=None, chunk_size=None, n_splits=None,
                                              n_jobs=n_jobs), min(4 * n_jobs, len(self.test_data_numpy)))

    def test_chunk_size_priority_over_n_splits(self):
        """
        chunk_size should have priority over n_splits
        """
        with self.subTest(input='list', chunk_size=1, n_splits=6):
            self.assertEqual(get_n_chunks(self.test_data, iterable_len=None, chunk_size=1, n_splits=6, n_jobs=None), 13)
        with self.subTest(input='numpy', chunk_size=1, n_splits=6):
            self.assertEqual(get_n_chunks(self.test_data_numpy, iterable_len=None, chunk_size=1, n_splits=6,
                                          n_jobs=None), 100)

        with self.subTest(input='list', chunk_size=3, n_splits=3):
            self.assertEqual(get_n_chunks(self.test_data, iterable_len=None, chunk_size=3, n_splits=3, n_jobs=None), 5)
        with self.subTest(input='numpy', chunk_size=3, n_splits=3):
            self.assertEqual(get_n_chunks(self.test_data_numpy, iterable_len=None, chunk_size=3, n_splits=3,
                                          n_jobs=None), 34)

    def test_generator_input_with_no_iterable_len_raises(self):
        """
        When working with generators the iterable_len should be provided (the working examples are already tested above)
        """
        for chunk_size, n_splits, n_jobs in product([None, 1, 3], [None, 1, 3], [None, 1, 3]):
            with self.subTest(chunk_size=chunk_size, n_splits=n_splits, n_jobs=n_jobs), self.assertRaises(ValueError):
                get_n_chunks(iter(self.test_data), iterable_len=None, chunk_size=chunk_size, n_splits=n_splits,
                             n_jobs=n_jobs)


class MakeSingleArgumentsTest(unittest.TestCase):

    def test_make_single_arguments(self):
        """
        Tests the make_single_arguments function for different inputs
        """
        # Test for some different inputs
        for (args_in, args_out), generator in product(
                [(['a', 'c', 'b', 'd'], [('a',), ('c',), ('b',), ('d',)]),
                 ([1, 2, 3, 4, 5], [(1,), (2,), (3,), (4,), (5,)]),
                 ([(True,), (False,), (None,)], [((True,),), ((False,),), ((None,),)])],
                [False, True]
        ):
            # Transform
            args_transformed = make_single_arguments((arg for arg in args_in) if generator else args_in,
                                                     generator=generator)

            # Check type
            self.assertTrue(isinstance(args_transformed, types.GeneratorType if generator else list))

            # Check contents
            self.assertEqual(list(args_transformed), args_out)


class FormatSecondsTest(unittest.TestCase):

    def test_none_input(self):
        """
        When the input is None it should return an empty string
        """
        for with_milliseconds in [False, True]:
            with self.subTest(with_milliseconds=with_milliseconds):
                self.assertEqual(format_seconds(None, with_milliseconds=with_milliseconds), '')

    def test_without_milliseconds(self):
        """
        Test output without milliseconds
        """
        for seconds, expected_output in [(0, '0:00:00'), (1, '0:00:01'), (1.337, '0:00:01'), (2.9, '0:00:02'),
                                         (123456.78901234, '1 day, 10:17:36')]:
            with self.subTest(seconds=seconds):
                self.assertEqual(format_seconds(seconds, with_milliseconds=False), expected_output)

    def test_with_milliseconds(self):
        """
        Test output with milliseconds. Only shows them when they're actually needed.
        """
        for seconds, expected_output in [(0, '0:00:00'), (1, '0:00:01'), (1.337, '0:00:01.337'), (2.9, '0:00:02.900'),
                                         (123456.78901234, '1 day, 10:17:36.789')]:
            with self.subTest(seconds=seconds):
                self.assertEqual(format_seconds(seconds, with_milliseconds=True), expected_output)


class TimeItTest(unittest.TestCase):

    def test_array_storage(self):
        """
        TimeIt should write to the correct idx in the cum_time_array container. The max_time_array is a min-heap
        container, so the lowest value is stored at index 0. The single highest value in this case is stored at index 2
        """
        for array_idx in range(5):
            cum_time_array = [0.0, 0.0, 0.0, 0.0, 0.0]
            max_time_array = [(0.0, ''), (0.0, ''), (0.0, ''), (0.0, ''), (0.0, '')]
            with self.subTest(array_idx=array_idx), patch('mpire.utils.time.time', side_effect=[0.0, 4.2]), \
                    TimeIt(cum_time_array, array_idx, max_time_array):
                pass
            self.assertListEqual([t for idx, t in enumerate(cum_time_array) if idx != array_idx], [0.0, 0.0, 0.0, 0.0])
            self.assertListEqual([t for idx, t in enumerate(max_time_array) if idx != 2.0],
                                 [(0.0, ''), (0.0, ''), (0.0, ''), (0.0, '')])
            self.assertEqual(cum_time_array[array_idx], 4.2)
            self.assertGreaterEqual(max_time_array[2], (4.2, None))

    def test_cum_time(self):
        """
        Using TimeIt multiple times should increase the cum_time_array
        """
        # These return values are used by TimeIt in order: start, end, start, end, ... So the first time the duration
        # will be 1 second, then 2 seconds, and 3 seconds.
        cum_time_array = [0]
        with patch('mpire.utils.time.time', side_effect=[0.0, 1.0, 0.0, 2.0, 0.0, 3.0]):
            with TimeIt(cum_time_array, 0):
                pass
            self.assertEqual(cum_time_array[0], 1.0)

            with TimeIt(cum_time_array, 0):
                pass
            self.assertEqual(cum_time_array[0], 3.0)

            with TimeIt(cum_time_array, 0):
                pass
            self.assertEqual(cum_time_array[0], 6.0)

    def test_max_time(self):
        """
        Using TimeIt multiple times should store the max duration value in the max_time_array using heapq. There's only
        room for the highest 5 values, while it is called 6 times. The smallest duration shouldn't be present.
        """
        # These return values are used by TimeIt in order: start, end, start, end, ... So the first time the duration
        # will be 1 second, then 2 seconds, 3 seconds, 3 seconds again, 0.5 seconds, and 10 seconds.
        cum_time_array = [0.0]
        max_time_array = [(0.0, ''), (0.0, ''), (0.0, ''), (0.0, ''), (0.0, '')]
        with patch('mpire.utils.time.time', side_effect=[0.0, 1.0, 0.0, 2.0, 0.0, 3.0, 0.0, 3.0, 0.0, 0.5, 0.0, 10.0]):
            for _ in range(6):
                with TimeIt(cum_time_array, 0, max_time_array):
                    pass
            self.assertListEqual(max_time_array, [(1.0, None), (2.0, None), (10.0, None), (3.0, None), (3.0, None)])

    def test_format_args(self):
        """
        The format args func should be called when provided
        """
        for format_func, formatted in [(lambda: "1", "1"), (lambda: 2, 2), (lambda: "foo", "foo")]:
            # These return values are used by TimeIt in order: start, end, start, end, ... So the first time the
            # duration will be 1 second, then 2 seconds, and 3 seconds.
            with self.subTest(format_func=format_func), \
                    patch('mpire.utils.time.time', side_effect=[0.0, 1.0, 0.0, 2.0, 0.0, 3.0]):
                cum_time_array = [0.0]
                max_time_array = [(0.0, ''), (0.0, '')]
                for _ in range(3):
                    with TimeIt(cum_time_array, 0, max_time_array, format_func):
                        pass

                # The heapq only had room for two entries. The highest durations should be kept
                self.assertListEqual(max_time_array, [(2.0, formatted), (3.0, formatted)])