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, 'functools.partial', 'partial', False, 'from functools import partial\n'), (49, 'rosetta.parallel.parallel_easy.map_easy', 'parallel_easy.map_easy', False, 'from rosetta.parallel import parallel_easy, pandas_easy\n'), (53, 'rosetta.parallel.parallel_easy.map_easy', 'parallel_easy.map_easy', False, 'from rosetta.parallel import parallel_easy, pandas_easy\n'), (57, 'rosetta.parallel.parallel_easy.imap_easy', 'parallel_easy.imap_easy', False, 'from rosetta.parallel import parallel_easy, pandas_easy\n'), (64, 'rosetta.parallel.parallel_easy.imap_easy', 'parallel_easy.imap_easy', False, 'from rosetta.parallel import parallel_easy, pandas_easy\n'), (92, 'numpy.random.randint', 'np.random.randint', True, 'import numpy as np\n'), (93, 'numpy.random.randint', 'np.random.randint', True, 'import numpy as np\n'), (94, 'numpy.random.randint', 'np.random.randint', True, 'import numpy as np\n'), (99, '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, '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, 'pandas.DataFrame', 'pd.DataFrame', True, 'import pandas as pd\n'), (159, '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, 'pandas.util.testing.assert_series_equal', 'assert_series_equal', False, 'from pandas.util.testing import assert_frame_equal, assert_series_equal\n'), (163, 'pandas.Series', 'pd.Series', True, 'import pandas as pd\n'), (166, '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, 'pandas.util.testing.assert_series_equal', 'assert_series_equal', False, 'from pandas.util.testing import assert_frame_equal, assert_series_equal\n'), (171, 'pandas.DataFrame', 'pd.DataFrame', True, 'import pandas as pd\n'), (174, '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, 'pandas.util.testing.assert_frame_equal', 'assert_frame_equal', False, 'from pandas.util.testing import assert_frame_equal, assert_series_equal\n'), (179, 'pandas.DataFrame', 'pd.DataFrame', True, 'import pandas as pd\n'), (182, '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, 'pandas.util.testing.assert_frame_equal', 'assert_frame_equal', False, 'from pandas.util.testing import assert_frame_equal, assert_series_equal\n'), (208, 'threading.Lock', 'threading.Lock', False, 'import threading\n'), (209, 'StringIO.StringIO', 'StringIO', False, 'from StringIO import StringIO\n'), (225, 'StringIO.StringIO', 'StringIO', False, 'from StringIO import StringIO\n'), (226, 'rosetta.parallel.threading_easy.threading_easy', 'threading_easy', False, 'from rosetta.parallel.threading_easy import threading_easy, LockIterateApply\n'), (233, 'StringIO.StringIO', 'StringIO', False, 'from StringIO import StringIO\n'), (234, 'rosetta.parallel.threading_easy.threading_easy', 'threading_easy', False, 'from rosetta.parallel.threading_easy import threading_easy, LockIterateApply\n'), (75, 'rosetta.parallel.parallel_easy._n_jobs_wrap', 'parallel_easy._n_jobs_wrap', False, 'from rosetta.parallel import parallel_easy, pandas_easy\n'), (211, 'rosetta.parallel.threading_easy.LockIterateApply', 'LockIterateApply', False, 'from rosetta.parallel.threading_easy import threading_easy, LockIterateApply\n')]