import unittest
from functools import partial

import pandas as pd
from pandas.util.testing import assert_frame_equal, assert_series_equal
import numpy as np
import threading
from StringIO import StringIO

from rosetta.parallel import parallel_easy, pandas_easy
from rosetta.parallel.threading_easy import threading_easy, LockIterateApply


# A couple functions for testing parallel easy
# Must be defined outside of the test class for some reason.
def _abfunc(x, a, b=1):
    return x * a * b
abfunc = partial(_abfunc, 2, 3)


def frame_to_series(frame):
    x = frame.iloc[0, 0]
    return pd.Series([x] * len(frame.columns), index=frame.columns)


def rightmax(mylist):
    return [max(mylist[i: i+2]) for i in range(len(mylist))]


def leftmax(mylist):
    for i in range(len(mylist)):
        if i == 0:
            result = [mylist[0]]
        else:
            result.append(max(mylist[i - 1: i+1]))

    return result


class TestBase(unittest.TestCase):
    """
    Tests the parallel_easy module.
    """
    def setUp(self):
        self.numbers = range(5)
        self.benchmark = [0, 6, 12, 18, 24]

    def test_map_easy_1job(self):
        result = parallel_easy.map_easy(abfunc, self.numbers, 1)
        self.assertEqual(result, self.benchmark)

    def test_map_easy_3job(self):
        result = parallel_easy.map_easy(abfunc, self.numbers, 3)
        self.assertEqual(result, self.benchmark)

    def test_imap_easy_1job(self):
        result_iterator = parallel_easy.imap_easy(abfunc, self.numbers, 1, 1)
        result = []
        for number in result_iterator:
            result.append(number)
        self.assertEqual(result, self.benchmark)

    def test_imap_easy_3job(self):
        result_iterator = parallel_easy.imap_easy(abfunc, self.numbers, 3, 1)
        result = []
        for number in result_iterator:
            result.append(number)
        self.assertEqual(result, self.benchmark)

    def test_n_jobs_wrap_positive(self):
        """
        For n_jobs positive, the wrap should return n_jobs.
        """
        for n_jobs in range(1, 5):
            result = parallel_easy._n_jobs_wrap(n_jobs)
            self.assertEqual(result, n_jobs)

    def test_n_jobs_wrap_zero(self):
        """
        For n_jobs zero, the wrap should raise a ValueError
        """
        self.assertRaises(ValueError, parallel_easy._n_jobs_wrap, 0)


class TestMapEasyPaddedBlock(unittest.TestCase):
    """
    Tests the parallel_easy.map_easy_padded_blocks function.
    """
    def setUp(self):
        #self.numbers_1 = [
        #    0, 0, 2, -1, 4, 2, 6, 7, 6, 9, 12, 11, 11, 14, 55, 55, 44, 33, 33]
        self.numbers_10 = np.random.randint(0, 5, 10)
        self.numbers_101 = np.random.randint(0, 5, 101)
        self.numbers_51 = np.random.randint(0, 5, 101)
        #self.numbers_1 = [0, 1, 2, 0, 3, 2, 4, 3, 2, 3, 3]
        self.n_jobs = 1

    def lefttest(self, numbers, buffer_len, blocksize):
        result = parallel_easy.map_easy_padded_blocks(
            leftmax, numbers, self.n_jobs, buffer_len, blocksize=blocksize)
        benchmark = leftmax(numbers)
        self.assertEqual(result, benchmark)

    def righttest(self, numbers, buffer_len, blocksize):
        result = parallel_easy.map_easy_padded_blocks(
            rightmax, numbers, self.n_jobs, buffer_len, blocksize=blocksize)
        benchmark = rightmax(numbers)
        self.assertEqual(result, benchmark)

    def test_map_easy_padded_blocks_14(self):
        buffer_len = 1
        blocksize = 4
        self.lefttest(self.numbers_10, buffer_len, blocksize)
        self.lefttest(self.numbers_101, buffer_len, blocksize)
        self.lefttest(self.numbers_51, buffer_len, blocksize)
        self.righttest(self.numbers_10, buffer_len, blocksize)
        self.righttest(self.numbers_101, buffer_len, blocksize)
        self.righttest(self.numbers_51, buffer_len, blocksize)

    def test_map_easy_padded_blocks_24(self):
        buffer_len = 2
        blocksize = 4
        self.lefttest(self.numbers_10, buffer_len, blocksize)
        self.lefttest(self.numbers_101, buffer_len, blocksize)
        self.lefttest(self.numbers_51, buffer_len, blocksize)
        self.righttest(self.numbers_10, buffer_len, blocksize)
        self.righttest(self.numbers_101, buffer_len, blocksize)
        self.righttest(self.numbers_51, buffer_len, blocksize)

    def test_map_easy_padded_blocks_37(self):
        buffer_len = 3
        blocksize = 7
        self.lefttest(self.numbers_101, buffer_len, blocksize)
        self.lefttest(self.numbers_51, buffer_len, blocksize)
        self.righttest(self.numbers_101, buffer_len, blocksize)
        self.righttest(self.numbers_51, buffer_len, blocksize)

    def test_map_easy_padded_blocks_17(self):
        buffer_len = 1
        blocksize = 7
        self.lefttest(self.numbers_10, buffer_len, blocksize)
        self.lefttest(self.numbers_101, buffer_len, blocksize)
        self.lefttest(self.numbers_51, buffer_len, blocksize)
        self.righttest(self.numbers_10, buffer_len, blocksize)
        self.righttest(self.numbers_101, buffer_len, blocksize)
        self.righttest(self.numbers_51, buffer_len, blocksize)


class TestPandasEasy(unittest.TestCase):
    """
    Tests the pandas_easy module.
    """
    def setUp(self):
        pass

    def test_groupby_to_scalar_to_series_1(self):
        df = pd.DataFrame({'a': [6, 2, 2], 'b': [4, 5, 6]})
        benchmark = df.groupby('a').apply(max)
        result = pandas_easy.groupby_to_scalar_to_series(df, max, 1, by='a')
        assert_series_equal(result, benchmark)

    def test_groupby_to_scalar_to_series_2(self):
        s = pd.Series([1, 2, 3, 4])
        labels = ['a', 'a', 'b', 'b']
        benchmark = s.groupby(labels).apply(max)
        result = pandas_easy.groupby_to_scalar_to_series(
            s, max, 1, by=labels)
        assert_series_equal(result, benchmark)

    def test_groupby_to_series_to_frame_1(self):
        df = pd.DataFrame({'a': [6, 2, 2], 'b': [4, 5, 6]})
        labels = ['g1', 'g1', 'g2']
        benchmark = df.groupby(labels).mean()
        result = pandas_easy.groupby_to_series_to_frame(
            df, np.mean, 1, use_apply=True, by=labels)
        assert_frame_equal(result, benchmark)

    def test_groupby_to_series_to_frame_2(self):
        df = pd.DataFrame({'a': [6, 2, 2], 'b': [4, 5, 6]})
        labels = ['g1', 'g1', 'g2']
        benchmark = df.groupby(labels).apply(frame_to_series)
        result = pandas_easy.groupby_to_series_to_frame(
            df, frame_to_series, 1, use_apply=False, by=labels)
        assert_frame_equal(result, benchmark)


class TestLockIterateApply(unittest.TestCase):
    """
    Test the Locked Iterator Class
    """
    def setUp(self):
        self.data = ['my', 'name', 'is', 'daniel']
        self.num_threads = 4

        def bytwo(x):
            return 2 * x

        self.func = bytwo

        def it():
            for i in self.data:
                yield i

        self.myiter = it()

    def test_locked_iterator(self):
        threads = []
        lock = threading.Lock()
        out = StringIO()
        for i in range(self.num_threads):
            t = LockIterateApply(self.func, self.myiter, lock, ',', out)
            threads.append(t)

        for t in threads:
            t.start()

        for t in threads:
            t.join()

        benchmark = set(['mymy', 'namename', 'danieldaniel', 'isis', ''])
        results = set(out.getvalue().split(','))
        self.assertEqual(results, benchmark)

    def test_threading_easy(self):
        out = StringIO()
        threading_easy(self.func, self.myiter, self.num_threads, ',', out)

        benchmark = set(['mymy', 'namename', 'danieldaniel', 'isis', ''])
        results = set(out.getvalue().split(','))
        self.assertEqual(results, benchmark)

    def test_threading_easy_single(self):
        out = StringIO()
        threading_easy(self.func, self.myiter, 1, ',', out)

        benchmark = set(['mymy', 'namename', 'danieldaniel', 'isis', ''])
        results = set(out.getvalue().split(','))
        self.assertEqual(results, benchmark)


[((18, 9, 18, 31), 'functools.partial', 'partial', ({(18, 17, 18, 24): '_abfunc', (18, 26, 18, 27): '2', (18, 29, 18, 30): '3'}, {}), '(_abfunc, 2, 3)', False, 'from functools import partial\n'), ((49, 17, 49, 64), 'rosetta.parallel.parallel_easy.map_easy', 'parallel_easy.map_easy', ({(49, 40, 49, 46): 'abfunc', (49, 48, 49, 60): 'self.numbers', (49, 62, 49, 63): '1'}, {}), '(abfunc, self.numbers, 1)', False, 'from rosetta.parallel import parallel_easy, pandas_easy\n'), ((53, 17, 53, 64), 'rosetta.parallel.parallel_easy.map_easy', 'parallel_easy.map_easy', ({(53, 40, 53, 46): 'abfunc', (53, 48, 53, 60): 'self.numbers', (53, 62, 53, 63): '3'}, {}), '(abfunc, self.numbers, 3)', False, 'from rosetta.parallel import parallel_easy, pandas_easy\n'), ((57, 26, 57, 77), 'rosetta.parallel.parallel_easy.imap_easy', 'parallel_easy.imap_easy', ({(57, 50, 57, 56): 'abfunc', (57, 58, 57, 70): 'self.numbers', (57, 72, 57, 73): '1', (57, 75, 57, 76): '1'}, {}), '(abfunc, self.numbers, 1, 1)', False, 'from rosetta.parallel import parallel_easy, pandas_easy\n'), ((64, 26, 64, 77), 'rosetta.parallel.parallel_easy.imap_easy', 'parallel_easy.imap_easy', ({(64, 50, 64, 56): 'abfunc', (64, 58, 64, 70): 'self.numbers', (64, 72, 64, 73): '3', (64, 75, 64, 76): '1'}, {}), '(abfunc, self.numbers, 3, 1)', False, 'from rosetta.parallel import parallel_easy, pandas_easy\n'), ((92, 26, 92, 53), 'numpy.random.randint', 'np.random.randint', ({(92, 44, 92, 45): '0', (92, 47, 92, 48): '5', (92, 50, 92, 52): '10'}, {}), '(0, 5, 10)', True, 'import numpy as np\n'), ((93, 27, 93, 55), 'numpy.random.randint', 'np.random.randint', ({(93, 45, 93, 46): '0', (93, 48, 93, 49): '5', (93, 51, 93, 54): '101'}, {}), '(0, 5, 101)', True, 'import numpy as np\n'), ((94, 26, 94, 54), 'numpy.random.randint', 'np.random.randint', ({(94, 44, 94, 45): '0', (94, 47, 94, 48): '5', (94, 50, 94, 53): '101'}, {}), '(0, 5, 101)', True, 'import numpy as np\n'), ((99, 17, 100, 75), 'rosetta.parallel.parallel_easy.map_easy_padded_blocks', 'parallel_easy.map_easy_padded_blocks', (), '', False, 'from rosetta.parallel import parallel_easy, pandas_easy\n'), ((105, 17, 106, 76), 'rosetta.parallel.parallel_easy.map_easy_padded_blocks', 'parallel_easy.map_easy_padded_blocks', (), '', False, 'from rosetta.parallel import parallel_easy, pandas_easy\n'), ((157, 13, 157, 59), 'pandas.DataFrame', 'pd.DataFrame', ({(157, 26, 157, 58): "{'a': [6, 2, 2], 'b': [4, 5, 6]}"}, {}), "({'a': [6, 2, 2], 'b': [4, 5, 6]})", True, 'import pandas as pd\n'), ((159, 17, 159, 76), 'rosetta.parallel.pandas_easy.groupby_to_scalar_to_series', 'pandas_easy.groupby_to_scalar_to_series', (), '', False, 'from rosetta.parallel import parallel_easy, pandas_easy\n'), ((160, 8, 160, 46), 'pandas.util.testing.assert_series_equal', 'assert_series_equal', ({(160, 28, 160, 34): 'result', (160, 36, 160, 45): 'benchmark'}, {}), '(result, benchmark)', False, 'from pandas.util.testing import assert_frame_equal, assert_series_equal\n'), ((163, 12, 163, 35), 'pandas.Series', 'pd.Series', ({(163, 22, 163, 34): '[1, 2, 3, 4]'}, {}), '([1, 2, 3, 4])', True, 'import pandas as pd\n'), ((166, 17, 167, 33), 'rosetta.parallel.pandas_easy.groupby_to_scalar_to_series', 'pandas_easy.groupby_to_scalar_to_series', (), '', False, 'from rosetta.parallel import parallel_easy, pandas_easy\n'), ((168, 8, 168, 46), 'pandas.util.testing.assert_series_equal', 'assert_series_equal', ({(168, 28, 168, 34): 'result', (168, 36, 168, 45): 'benchmark'}, {}), '(result, benchmark)', False, 'from pandas.util.testing import assert_frame_equal, assert_series_equal\n'), ((171, 13, 171, 59), 'pandas.DataFrame', 'pd.DataFrame', ({(171, 26, 171, 58): "{'a': [6, 2, 2], 'b': [4, 5, 6]}"}, {}), "({'a': [6, 2, 2], 'b': [4, 5, 6]})", True, 'import pandas as pd\n'), ((174, 17, 175, 54), 'rosetta.parallel.pandas_easy.groupby_to_series_to_frame', 'pandas_easy.groupby_to_series_to_frame', (), '', False, 'from rosetta.parallel import parallel_easy, pandas_easy\n'), ((176, 8, 176, 45), 'pandas.util.testing.assert_frame_equal', 'assert_frame_equal', ({(176, 27, 176, 33): 'result', (176, 35, 176, 44): 'benchmark'}, {}), '(result, benchmark)', False, 'from pandas.util.testing import assert_frame_equal, assert_series_equal\n'), ((179, 13, 179, 59), 'pandas.DataFrame', 'pd.DataFrame', ({(179, 26, 179, 58): "{'a': [6, 2, 2], 'b': [4, 5, 6]}"}, {}), "({'a': [6, 2, 2], 'b': [4, 5, 6]})", True, 'import pandas as pd\n'), ((182, 17, 183, 63), 'rosetta.parallel.pandas_easy.groupby_to_series_to_frame', 'pandas_easy.groupby_to_series_to_frame', (), '', False, 'from rosetta.parallel import parallel_easy, pandas_easy\n'), ((184, 8, 184, 45), 'pandas.util.testing.assert_frame_equal', 'assert_frame_equal', ({(184, 27, 184, 33): 'result', (184, 35, 184, 44): 'benchmark'}, {}), '(result, benchmark)', False, 'from pandas.util.testing import assert_frame_equal, assert_series_equal\n'), ((208, 15, 208, 31), 'threading.Lock', 'threading.Lock', ({}, {}), '()', False, 'import threading\n'), ((209, 14, 209, 24), 'StringIO.StringIO', 'StringIO', ({}, {}), '()', False, 'from StringIO import StringIO\n'), ((225, 14, 225, 24), 'StringIO.StringIO', 'StringIO', ({}, {}), '()', False, 'from StringIO import StringIO\n'), ((226, 8, 226, 74), 'rosetta.parallel.threading_easy.threading_easy', 'threading_easy', ({(226, 23, 226, 32): 'self.func', (226, 34, 226, 45): 'self.myiter', (226, 47, 226, 63): 'self.num_threads', (226, 65, 226, 68): '""","""', (226, 70, 226, 73): 'out'}, {}), "(self.func, self.myiter, self.num_threads, ',', out)", False, 'from rosetta.parallel.threading_easy import threading_easy, LockIterateApply\n'), ((233, 14, 233, 24), 'StringIO.StringIO', 'StringIO', ({}, {}), '()', False, 'from StringIO import StringIO\n'), ((234, 8, 234, 59), 'rosetta.parallel.threading_easy.threading_easy', 'threading_easy', ({(234, 23, 234, 32): 'self.func', (234, 34, 234, 45): 'self.myiter', (234, 47, 234, 48): '(1)', (234, 50, 234, 53): '""","""', (234, 55, 234, 58): 'out'}, {}), "(self.func, self.myiter, 1, ',', out)", False, 'from rosetta.parallel.threading_easy import threading_easy, LockIterateApply\n'), ((75, 21, 75, 55), 'rosetta.parallel.parallel_easy._n_jobs_wrap', 'parallel_easy._n_jobs_wrap', ({(75, 48, 75, 54): 'n_jobs'}, {}), '(n_jobs)', False, 'from rosetta.parallel import parallel_easy, pandas_easy\n'), ((211, 16, 211, 72), 'rosetta.parallel.threading_easy.LockIterateApply', 'LockIterateApply', ({(211, 33, 211, 42): 'self.func', (211, 44, 211, 55): 'self.myiter', (211, 57, 211, 61): 'lock', (211, 63, 211, 66): '""","""', (211, 68, 211, 71): 'out'}, {}), "(self.func, self.myiter, lock, ',', out)", False, 'from rosetta.parallel.threading_easy import threading_easy, LockIterateApply\n')]