import pickle
import sys
import time
import unittest
from io import BytesIO
from unittest.mock import MagicMock

from pyxtension.Json import JsonList, Json

from pyxtension.streams import stream, slist, sset, sdict, ItrFromFunc, defaultstreamdict

ifilter = filter
xrange = range

__author__ = 'ASU'


# s = stream((1, 2, 3))
# a = s.map(lambda _: int(_))[1:2][0]

class StreamTestCase(unittest.TestCase):
    def setUp(self):
        self.s = lambda: stream((1, 2, 3))
    
    def testStream(self):
        s = self.s
        self.assertEquals(list(ifilter(lambda i: i % 2 == 0, s())), [2])
        self.assertEquals(list(s().filter(lambda i: i % 2 == 0)), [2])
        self.assertEquals(s().filter(lambda i: i % 2 == 0).toList(), [2])
        self.assertEquals(s()[1], 2)
        self.assertEquals(s()[1:].toList(), [2, 3])
        self.assertEqual(s().take(2).toList(), [1, 2])
        self.assertAlmostEqual(stream((0, 1, 2, 3)).filter(lambda x: x > 0).entropy(), 1.4591479)
        self.assertEquals(stream([(1, 2), (3, 4)]).zip().toList(), [(1, 3), (2, 4)])
    
    def test_filterFromGeneratorReinstantiatesProperly(self):
        s = stream(ItrFromFunc(lambda: (i for i in xrange(5))))
        s = s.filter(lambda e: e % 2 == 0)
        self.assertEquals(s.toList(), [0, 2, 4])
        self.assertEquals(s.toList(), [0, 2, 4])
        s = stream(xrange(5)).filter(lambda e: e % 2 == 0)
        self.assertEquals(s.toList(), [0, 2, 4])
        self.assertEquals(s.toList(), [0, 2, 4])
    
    def test_streamExists(self):
        s = stream([0, 1])
        self.assertEqual(s.exists(lambda e: e == 0), True)
        self.assertEqual(s.exists(lambda e: e == 2), False)
    
    def test_stream_str_doesntChangeStream(self):
        s = stream(iter((1, 2, 3, 4)))
        str(s)
        self.assertListEqual(s.toList(), [1, 2, 3, 4])
    
    def test_stream_repr_doesntChangeStream(self):
        s = stream(iter((1, 2, 3, 4)))
        repr(s)
        self.assertListEqual(s.toList(), [1, 2, 3, 4])
    
    def test_slist_str_nominal(self):
        l = [1, 2, 3]
        s = slist(l)
        s1 = str(s)
        self.assertEquals(str(s), str(l))
    
    def test_slist_repr_nominal(self):
        l = [1, 2, 3]
        s = slist(l)
        self.assertEquals(repr(s), repr(l))
    
    def testStreamToJson(self):
        j = stream((("a", 2), (3, 4))).toJson()
        self.assertIsInstance(j, JsonList)
        self.assertListEqual(j, [["a", 2], [3, 4]])
    
    def testSdictToJson(self):
        j = stream((("a", 2), (3, 4))).toMap().toJson()
        self.assertIsInstance(j, Json)
        self.assertEqual(j.a, 2)
        self.assertDictEqual(j, {'a': 2, 3: 4})
    
    def testStreamList(self):
        l = lambda: slist((1, 2, 3))
        self.assertEqual(l().toList(), [1, 2, 3])
        self.assertEqual(l()[-1], 3)
    
    def testStreamSet(self):
        s = lambda: sset([1, 2, 3, 2])
        self.assertEqual(s().size(), 3)
        self.assertEqual(s().map(lambda x: x).toList(), [1, 2, 3])
        self.assertEqual(len(s()), 3)
    
    def test_sdict(self):
        d = sdict({1: 2, 3: 4})
        self.assertListEqual(d.items().map(lambda t: t).toList(), [(1, 2), (3, 4)])
    
    def testStreamsFromGenerator(self):
        sg = stream(ItrFromFunc(lambda: (i for i in range(4))))
        self.assertEqual(sg.size(), 4)
        self.assertEqual(sg.size(), 4)
        self.assertEqual(sg.filter(lambda x: x > 1).toList(), [2, 3])
        self.assertEqual(sg.filter(lambda x: x > 1).toList(), [2, 3])
        self.assertEqual(sg.map(lambda x: x > 1).toList(), [False, False, True, True])
        self.assertEqual(sg.map(lambda x: x > 1).toList(), [False, False, True, True])
        self.assertEqual(sg.head(), 0)
        self.assertEqual(sg.head(), 0)
        self.assertEqual(sg.map(lambda i: i ** 2).enumerate().toList(), [(0, 0), (1, 1), (2, 4), (3, 9)])
        self.assertEqual(sg.reduce(lambda x, y: x + y, 5), 11)
    
    def testStreamPickling(self):
        sio = BytesIO()
        expected = slist(slist((i,)) for i in xrange(10))
        expected.dumpToPickle(sio)
        sio = BytesIO(sio.getvalue())
        
        result = stream.loadFromPickled(sio)
        self.assertListEqual(expected, list(result))
    
    def test_StreamFileReading(self):
        sio = BytesIO()
        expected = slist(slist((i,)) for i in xrange(10))
        expected.dumpToPickle(sio)
        sio = BytesIO(sio.getvalue())
        
        result = stream.loadFromPickled(sio)
        self.assertEquals(list(expected), list(result))
    
    def test_flatMap_nominal(self):
        s = stream([[1, 2], [3, 4], [4, 5]])
        self.assertListEqual(s.flatMap().toList(), [1, 2, 3, 4, 4, 5])
    
    def test_flatMap_withPredicate(self):
        s = stream(({1: 2, 3: 4}, {5: 6, 7: 8}))
        self.assertEquals(s.flatMap(dict.items).toSet(), set(((1, 2), (5, 6), (3, 4), (7, 8))))
    
    def test_flatMap_reiteration(self):
        l = stream(ItrFromFunc(lambda: (xrange(i) for i in xrange(5)))).flatMap()
        self.assertEquals(l.toList(), [0, 0, 1, 0, 1, 2, 0, 1, 2, 3])
        self.assertEquals(l.toList(),
                          [0, 0, 1, 0, 1, 2, 0, 1, 2, 3])  # second time to assert the regeneration of generator
    
    def test_flatMap_defaultIdentityFunction(self):
        l = slist(({1: 2, 3: 4}, {5: 6, 7: 8}))
        self.assertEquals(l.flatMap().toSet(), set((1, 3, 5, 7)))
    
    def test_sset_updateReturnsSelf(self):
        s = sset((1, 2))
        l = s.update((2, 3))
        self.assertEquals(l, set((1, 2, 3)))
    
    def test_sset_intersection_updateReturnsSelf(self):
        self.assertEquals(sset((1, 2)).update(set((2, 3))), set((1, 2, 3)))
    
    def test_reduceUsesInitProperly(self):
        self.assertEquals(slist([sset((1, 2)), sset((3, 4))]).reduce(lambda x, y: x.update(y)), set((1, 2, 3, 4)))
        self.assertEquals(slist([sset((1, 2)), sset((3, 4))]).reduce(lambda x, y: x.update(y), sset()),
                          set((1, 2, 3, 4)))
    
    def test_ssetChaining(self):
        s = sset().add(0).clear().add(1).add(2).remove(2).discard(3).update(set((3, 4, 5))) \
            .intersection_update(set((1, 3, 4))).difference_update(set((4,))).symmetric_difference_update(set((3, 4)))
        self.assertEquals(s, set((1, 4)))
    
    def test_maxes(self):
        self.assertEquals(stream(['a', 'abc', 'abcd', 'defg', 'cde']).maxes(lambda s: len(s)), ['abcd', 'defg'])
    
    def test_mins(self):
        self.assertEquals(stream(['abc', 'a', 'abcd', 'defg', 'cde']).mins(lambda s: len(s)), ['a'])
    
    def test_min_nominal(self):
        self.assertEqual(stream([2, 1]).min(), 1)
        self.assertEqual(stream(['abc', 'a']).min(key=len), 'a')
    
    def test_min_raises_on_empty_sequence(self):
        with self.assertRaises(ValueError):
            stream().min()
    
    def test_min_default_nominal(self):
        self.assertEqual(stream([2, 1]).min_default('default'), 1)
        self.assertEqual(stream(['abc', 'a']).min_default('default', key=len), 'a')
        self.assertEqual(stream().min_default('default'), 'default')
    
    def test_defaultstreamdictBasics(self):
        dd = defaultstreamdict(slist)
        dd[1].append(2)
        self.assertEquals(dd, {1: [2]})
    
    def test_defaultstreamdictSerialization(self):
        dd = defaultstreamdict(slist)
        dd[1].append(2)
        s = pickle.dumps(dd)
        newDd = pickle.loads(s)
        self.assertEquals(newDd, dd)
        self.assertIsInstance(newDd[1], slist)
    
    def test_stream_add(self):
        s1 = stream([1, 2])
        s2 = stream([3, 4])
        s3 = s1 + s2
        ll = s3.toList()
        self.assertEquals(s3.toList(), [1, 2, 3, 4])
        self.assertEquals(s3.toList(), [1, 2, 3, 4])  # second time to exclude one time iterator bug
        s1 = s1 + s2
        self.assertEquals(s1.toList(), [1, 2, 3, 4])
        self.assertEquals(s1.toList(), [1, 2, 3, 4])  # second time to exclude one time iterator bug
    
    def test_stream_iadd(self):
        s1 = stream([1, 2])
        s1 += [3, 4]
        s1 += stream(xrange(5, 6))  # use xrange to cover the iterator case
        self.assertEquals(s1.toList(), [1, 2, 3, 4, 5])
        self.assertEquals(s1.toList(), [1, 2, 3, 4, 5])  # second time to exclude one time iterator bug
        self.assertEquals(s1.toList(), [1, 2, 3, 4, 5])
    
    def test_stream_getitem(self):
        s = stream(i for i in xrange(1))
        self.assertEqual(s[0], 0)
    
    def test_fastmap_time(self):
        def sleepFunc(el):
            time.sleep(0.3)
            return el * el
        
        s = stream(xrange(100))
        t1 = time.time()
        res = s.fastmap(sleepFunc, poolSize=50).toSet()
        dt = time.time() - t1
        expected = set(i * i for i in xrange(100))
        self.assertSetEqual(res, expected)
        self.assertLessEqual(dt, 1.5)
    
    def test_fastmap_nominal(self):
        s = stream(xrange(100))
        res = s.fastmap(lambda x: x * x, poolSize=4).toSet()
        expected = set(i * i for i in xrange(100))
        self.assertSetEqual(res, expected)
    
    def test_fastmap_one_el(self):
        s = stream([1, ])
        res = s.fastmap(lambda x: x * x, poolSize=4).toSet()
        expected = set((1,))
        self.assertSetEqual(res, expected)
    
    def test_fastmap_no_el(self):
        s = stream([])
        res = s.fastmap(lambda x: x * x, poolSize=4).toSet()
        expected = set()
        self.assertSetEqual(res, expected)
    
    def test_fastmap_None_el(self):
        s = stream([None])
        res = s.fastmap(lambda x: x, poolSize=4).toSet()
        expected = set([None])
        self.assertSetEqual(res, expected)
    
    def test_fastmap_raises_exception(self):
        s = stream([None])
        with self.assertRaises(TypeError):
            res = s.fastmap(lambda x: x * x, poolSize=4).toSet()
    
    def test_unique_nominal(self):
        s = stream([1, 2, 3, 1, 2])
        self.assertListEqual(s.unique().toList(), [1, 2, 3])
    
    def test_unique_mapping(self):
        s = stream(['abc', 'def', 'a', 'b', 'ab'])
        self.assertListEqual(s.unique(len).toList(), ['abc', 'a', 'ab'])
    
    def test_unique_empty_stream(self):
        s = stream([])
        self.assertListEqual(s.unique().toList(), [])
    
    def test_unique_generator_stream(self):
        s = stream(ItrFromFunc(lambda: xrange(4)))
        u = s.unique()
        self.assertListEqual(u.toList(), [0, 1, 2, 3])
        self.assertListEqual(u.toList(), [0, 1, 2, 3])
    
    def test_pstddev_nominal(self):
        s = stream([1, 2, 3, 4])
        self.assertAlmostEqual(s.pstddev(), 1.118033988749895)
    
    def test_pstddev_exception(self):
        with self.assertRaises(ValueError):
            stream([]).pstddev()
    
    def test_mean(self):
        self.assertAlmostEqual(stream([1, 2, 3, 4]).mean(), 2.5)
    
    def test_mean_exception(self):
        with self.assertRaises(ValueError):
            stream([]).mean()
    
    def test_toSumCounter_nominal(self):
        s = stream([('a', 2), ('a', 4), ('b', 2.1), ('b', 3), ('c', 2)])
        self.assertDictEqual(s.toSumCounter(), {'a': 6, 'b': 5.1, 'c': 2})
    
    def test_toSumCounter_onEmptyStream(self):
        s = stream([])
        self.assertDictEqual(s.toSumCounter(), {})
    
    def test_toSumCounter_onStrings(self):
        s = stream([('a', 'b'), ('a', 'c')])
        self.assertDictEqual(s.toSumCounter(), {'a': 'bc'})
    
    def test_keyBy_nominal(self):
        self.assertListEqual(stream(['a', 'bb', '']).keyBy(len).toList(), [(1, 'a'), (2, 'bb'), (0, '')])
    
    def test_keys_nominal(self):
        self.assertListEqual(stream([(1, 'a'), (2, 'bb'), (0, '')]).keystream().toList(), [1, 2, 0])
    
    def test_values_nominal(self):
        self.assertListEqual(stream([(1, 'a'), (2, 'bb'), (0, '')]).values().toList(), ['a', 'bb', ''])
    
    def test_toMap(self):
        self.assertDictEqual(stream(((1, 2), (3, 4))).toMap(), {1: 2, 3: 4})
    
    def test_joinWithString(self):
        s = "|"
        strings = ('a', 'b', 'c')
        self.assertEqual(stream(iter(strings)).join(s), s.join(strings))
    
    def test_joinWithNone(self):
        s = ""
        strings = ('a', 'b', 'c')
        self.assertEqual(stream(iter(strings)).join(), s.join(strings))
    
    def test_joinWithFunction(self):
        class F:
            def __init__(self):
                self.counter = 0
            
            def __call__(self, *args, **kwargs):
                self.counter += 1
                return str(self.counter)
        
        strings = ('a', 'b', 'c')
        f = F()
        self.assertEqual(stream(iter(strings)).join(f), "a1b2c")
    
    def test_mkString(self):
        streamToTest = stream(('a', 'b', 'c'))
        mock = MagicMock()
        joiner = ","
        streamToTest.join = mock
        streamToTest.mkString(joiner)
        mock.assert_called_once_with(joiner)
    
    def test_reversedNominal(self):
        s = slist([1, 2, 3])
        self.assertListEqual(s.reversed().toList(), [3, 2, 1])
    
    def test_reversedException(self):
        s = stream(xrange(1, 2, 3))
        with self.assertRaises(TypeError):
            s.reversed()


"""
Allow for these test cases to be run from the command line
"""
if __name__ == '__main__':
    all_tests = unittest.TestLoader().loadTestsFromModule(sys.modules[__name__])
    unittest.TextTestRunner(verbosity=2).run(all_tests)
