import unittest

import rx
from rx import operators as ops
from rx.testing import TestScheduler, ReactiveTest

on_next = ReactiveTest.on_next
on_completed = ReactiveTest.on_completed
on_error = ReactiveTest.on_error
subscribe = ReactiveTest.subscribe
subscribed = ReactiveTest.subscribed
disposed = ReactiveTest.disposed
created = ReactiveTest.created


class RxException(Exception):
    pass

# Helper function for raising exceptions within lambdas


def _raise(ex):
    raise RxException(ex)


class TestGroupBy(unittest.TestCase):
    def test_group_by_with_key_comparer(self):
        scheduler = TestScheduler()
        key_invoked = [0]
        xs = scheduler.create_hot_observable(
            on_next(90, "error"),
            on_next(110, "error"),
            on_next(130, "error"),
            on_next(220, "  foo"),
            on_next(240, " FoO "),
            on_next(270, "baR  "),
            on_next(310, "foO "),
            on_next(350, " Baz   "),
            on_next(360, "  qux "),
            on_next(390, "   bar"),
            on_next(420, " BAR  "),
            on_next(470, "FOO "),
            on_next(480, "baz  "),
            on_next(510, " bAZ "),
            on_next(530, "    fOo    "),
            on_completed(570),
            on_next(580, "error"),
            on_completed(600),
            on_error(650, 'ex'))

        def factory():
            def key_mapper(x):
                key_invoked[0] += 1
                return x.lower().strip()

            return xs.pipe(
                    ops.group_by(key_mapper, lambda x: x),
                    ops.map(lambda g: g.key),
                    )

        results = scheduler.start(factory)
        assert results.messages == [
            on_next(220, "foo"),
            on_next(270, "bar"),
            on_next(350, "baz"),
            on_next(360, "qux"),
            on_completed(570)]
        assert xs.subscriptions == [subscribe(200, 570)]
        assert(key_invoked[0] == 12)

    def test_groupby_outer_complete(self):
        scheduler = TestScheduler()
        key_invoked = [0]
        ele_invoked = [0]
        xs = scheduler.create_hot_observable(
            on_next(90, "error"),
            on_next(110, "error"),
            on_next(130, "error"),
            on_next(220, "  foo"),
            on_next(240, " FoO "),
            on_next(270, "baR  "),
            on_next(310, "foO "),
            on_next(350, " Baz   "),
            on_next(360, "  qux "),
            on_next(390, "   bar"),
            on_next(420, " BAR  "),
            on_next(470, "FOO "),
            on_next(480, "baz  "),
            on_next(510, " bAZ "),
            on_next(530, "    fOo    "),
            on_completed(570),
            on_next(580, "error"),
            on_completed(600),
            on_error(650, 'ex'))

        def factory():
            def key_mapper(x):
                key_invoked[0] += 1
                return x.lower().strip()

            def element_mapper(x):
                ele_invoked[0] += 1
                return x[::-1]  # Yes, this is reverse string in Python

            return xs.pipe(
                    ops.group_by(key_mapper, element_mapper),
                    ops.map(lambda g: g.key),
                    )

        results = scheduler.start(factory)
        assert results.messages == [
            on_next(220, "foo"),
            on_next(270, "bar"),
            on_next(350, "baz"),
            on_next(360, "qux"),
            on_completed(570)]
        assert xs.subscriptions == [subscribe(200, 570)]
        assert key_invoked[0] == 12
        assert ele_invoked[0] == 12

    def test_group_by_outer_error(self):
        scheduler = TestScheduler()
        key_invoked = [0]
        ele_invoked = [0]
        ex = 'ex'
        xs = scheduler.create_hot_observable(
            on_next(90, "error"),
            on_next(110, "error"),
            on_next(130, "error"),
            on_next(220, "  foo"),
            on_next(240, " FoO "),
            on_next(270, "baR  "),
            on_next(310, "foO "),
            on_next(350, " Baz   "),
            on_next(360, "  qux "),
            on_next(390, "   bar"),
            on_next(420, " BAR  "),
            on_next(470, "FOO "),
            on_next(480, "baz  "),
            on_next(510, " bAZ "),
            on_next(530, "    fOo    "),
            on_error(570, ex),
            on_next(580, "error"),
            on_completed(600),
            on_error(650, 'ex'))

        def factory():
            def key_mapper(x):
                key_invoked[0] += 1
                return x.lower().strip()

            def element_mapper(x):
                ele_invoked[0] += 1
                return x[::-1]

            return xs.pipe(
                    ops.group_by(key_mapper, element_mapper),
                    ops.map(lambda g: g.key),
                    )

        results = scheduler.start(factory)

        assert results.messages == [
            on_next(220, "foo"),
            on_next(270, "bar"),
            on_next(350, "baz"),
            on_next(360, "qux"),
            on_error(570, ex)]
        assert xs.subscriptions == [subscribe(200, 570)]
        assert(key_invoked[0] == 12)
        assert(ele_invoked[0] == 12)

    def test_group_by_outer_dispose(self):
        scheduler = TestScheduler()
        key_invoked = [0]
        ele_invoked = [0]
        xs = scheduler.create_hot_observable(
            on_next(90, "error"),
            on_next(110, "error"),
            on_next(130, "error"),
            on_next(220, "  foo"),
            on_next(240, " FoO "),
            on_next(270, "baR  "),
            on_next(310, "foO "),
            on_next(350, " Baz   "),
            on_next(360, "  qux "),
            on_next(390, "   bar"),
            on_next(420, " BAR  "),
            on_next(470, "FOO "),
            on_next(480, "baz  "),
            on_next(510, " bAZ "),
            on_next(530, "    fOo    "),
            on_completed(570),
            on_next(580, "error"),
            on_completed(600),
            on_error(650, 'ex'))

        def factory():
            def key_mapper(x):
                key_invoked[0] += 1
                return x.lower().strip()

            def element_mapper(x):
                ele_invoked[0] += 1
                return x[::-1]

            return xs.pipe(
                    ops.group_by(key_mapper, element_mapper),
                    ops.map(lambda g: g.key),
                    )

        results = scheduler.start(factory, disposed=355)

        assert results.messages == [
            on_next(220, "foo"),
            on_next(270, "bar"),
            on_next(350, "baz")]
        assert xs.subscriptions == [subscribe(200, 355)]
        assert(key_invoked[0] == 5)
        assert(ele_invoked[0] == 5)

    def test_group_by_outer_key_on_error(self):
        scheduler = TestScheduler()
        key_invoked = [0]
        ele_invoked = [0]
        ex = 'ex'
        xs = scheduler.create_hot_observable(
            on_next(90, "error"),
            on_next(110, "error"),
            on_next(130, "error"),
            on_next(220, "  foo"),
            on_next(240, " FoO "),
            on_next(270, "baR  "),
            on_next(310, "foO "),
            on_next(350, " Baz   "),
            on_next(360, "  qux "),
            on_next(390, "   bar"),
            on_next(420, " BAR  "),
            on_next(470, "FOO "),
            on_next(480, "baz  "),
            on_next(510, " bAZ "),
            on_next(530, "    fOo    "),
            on_completed(570),
            on_next(580, "error"),
            on_completed(600),
            on_error(650, 'ex'))

        def factory():
            def key_mapper(x):
                key_invoked[0] += 1
                if key_invoked[0] == 10:
                    raise Exception(ex)

                return x.lower().strip()

            def element_mapper(x):
                ele_invoked[0] += 1
                return x[::-1]

            return xs.pipe(
                    ops.group_by(key_mapper, element_mapper),
                    ops.map(lambda g: g.key),
                    )

        results = scheduler.start(factory)
        assert results.messages == [
            on_next(220, "foo"),
            on_next(270, "bar"),
            on_next(350, "baz"),
            on_next(360, "qux"),
            on_error(480, ex)]
        assert xs.subscriptions == [subscribe(200, 480)]
        assert key_invoked[0] == 10
        assert ele_invoked[0] == 9

    def test_group_by_outer_ele_on_error(self):
        scheduler = TestScheduler()
        key_invoked = [0]
        ele_invoked = [0]
        ex = 'ex'
        xs = scheduler.create_hot_observable(
            on_next(90, "error"),
            on_next(110, "error"),
            on_next(130, "error"),
            on_next(220, "  foo"),
            on_next(240, " FoO "),
            on_next(270, "baR  "),
            on_next(310, "foO "),
            on_next(350, " Baz   "),
            on_next(360, "  qux "),
            on_next(390, "   bar"),
            on_next(420, " BAR  "),
            on_next(470, "FOO "),
            on_next(480, "baz  "),
            on_next(510, " bAZ "),
            on_next(530, "    fOo    "),
            on_completed(570),
            on_next(580, "error"),
            on_completed(600),
            on_error(650, 'ex'))

        def factory():
            def key_mapper(x):
                key_invoked[0] += 1
                return x.lower().strip()

            def element_mapper(x):
                ele_invoked[0] += 1
                if ele_invoked[0] == 10:
                    raise Exception(ex)
                return x[::-1]

            return xs.pipe(
                    ops.group_by(key_mapper, element_mapper),
                    ops.map(lambda g: g.key),
                    )

        results = scheduler.start(factory)
        assert results.messages == [
            on_next(220, "foo"),
            on_next(270, "bar"),
            on_next(350, "baz"),
            on_next(360, "qux"),
            on_error(480, ex)]
        assert xs.subscriptions == [subscribe(200, 480)]
        assert key_invoked[0] == 10
        assert ele_invoked[0] == 10

    def test_group_by_inner_complete(self):
        scheduler = TestScheduler()
        xs = scheduler.create_hot_observable(
            on_next(90, "error"),
            on_next(110, "error"),
            on_next(130, "error"),
            on_next(220, "  foo"),
            on_next(240, " FoO "),
            on_next(270, "baR  "),
            on_next(310, "foO "),
            on_next(350, " Baz   "),
            on_next(360, "  qux "),
            on_next(390, "   bar"),
            on_next(420, " BAR  "),
            on_next(470, "FOO "),
            on_next(480, "baz  "),
            on_next(510, " bAZ "),
            on_next(530, "    fOo    "),
            on_completed(570),
            on_next(580, "error"),
            on_completed(600),
            on_error(650, 'ex'))
        c = {
            "outer_subscription": None,
            "inner_subscriptions": {},
            "inners": {},
            "results": {},
            "outer": None
        }

        def action1(scheduler, state):
            c["outer"] = xs.pipe(
                    ops.group_by(lambda x: x.lower().strip(), lambda x: x[::-1]),
                    )

        scheduler.schedule_absolute(created, action1)

        def action2(scheduler, state):

            def next(group):

                result = scheduler.create_observer()
                c["inners"][group.key] = group
                c["results"][group.key] = result

                def action21(scheduler, state):
                    c["inner_subscriptions"][group.key] = group.subscribe(result, scheduler)

                scheduler.schedule_relative(100, action21)
            c["outer_subscription"] = c["outer"].subscribe(next, scheduler=scheduler)
        scheduler.schedule_absolute(subscribed, action2)

        def action3(scheduler, state):
            c["outer_subscription"].dispose()
            for sub in c["inner_subscriptions"].values():
                sub.dispose()

        scheduler.schedule_absolute(disposed, action3)
        scheduler.start()
        assert(len(c["inners"]) == 4)
        assert c["results"]['foo'].messages == [
            on_next(470, " OOF"),
            on_next(530, "    oOf    "),
            on_completed(570)]
        assert c["results"]['bar'].messages == [
            on_next(390, "rab   "),
            on_next(420, "  RAB "),
            on_completed(570)]
        assert c["results"]['baz'].messages == [
            on_next(480, "  zab"),
            on_next(510, " ZAb "),
            on_completed(570)]
        assert c["results"]['qux'].messages == [
            on_completed(570)]
        assert xs.subscriptions == [
            subscribe(200, 570)]

    def test_group_by_inner_complete_all(self):
        scheduler = TestScheduler()
        xs = scheduler.create_hot_observable(
            on_next(90, "error"),
            on_next(110, "error"),
            on_next(130, "error"),
            on_next(220, "  foo"),
            on_next(240, " FoO "),
            on_next(270, "baR  "),
            on_next(310, "foO "),
            on_next(350, " Baz   "),
            on_next(360, "  qux "),
            on_next(390, "   bar"),
            on_next(420, " BAR  "),
            on_next(470, "FOO "),
            on_next(480, "baz  "),
            on_next(510, " bAZ "),
            on_next(530, "    fOo    "),
            on_completed(570),
            on_next(580, "error"),
            on_completed(600),
            on_error(650, 'ex'))
        inners = {}
        inner_subscriptions = {}
        results = {}
        c = {
            "outer": None,
            "outer_subscription": None,
            "result": None
        }

        def action1(scheduler, state):
            c["outer"] = xs.pipe(
                    ops.group_by(
                            lambda x: x.lower().strip(),
                            lambda x: x[::-1],
                            ))

            return c["outer"]
        scheduler.schedule_absolute(created, action1)

        def action2(scheduler, state):
            def on_next(group):
                c["result"] = scheduler.create_observer()
                inners[group.key] = group
                results[group.key] = c["result"]
                inner_subscriptions[group.key] = group.subscribe(c["result"], scheduler)
            c["outer_subscription"] = c["outer"].subscribe(on_next, scheduler=scheduler)
            return c["outer_subscription"]
        scheduler.schedule_absolute(subscribed, action2)

        def action3(scheduler, state):
            c["outer_subscription"].dispose()
            for sub in inner_subscriptions.values():
                sub.dispose()
        scheduler.schedule_absolute(disposed, action3)

        scheduler.start()
        assert(len(inners) == 4)
        assert results['foo'].messages == [
            on_next(220, "oof  "),
            on_next(240, " OoF "),
            on_next(310, " Oof"),
            on_next(470, " OOF"),
            on_next(530, "    oOf    "),
            on_completed(570)]
        assert results['bar'].messages == [
            on_next(270, "  Rab"),
            on_next(390, "rab   "),
            on_next(420, "  RAB "),
            on_completed(570)]
        assert results['baz'].messages == [
            on_next(350, "   zaB "),
            on_next(480, "  zab"),
            on_next(510, " ZAb "),
            on_completed(570)]
        assert results['qux'].messages == [
            on_next(360, " xuq  "),
            on_completed(570)]
        assert xs.subscriptions == [
            subscribe(200, 570)]

    def test_group_by_inner_error(self):
        ex = 'ex1'
        scheduler = TestScheduler()
        xs = scheduler.create_hot_observable(
            on_next(90, "error"),
            on_next(110, "error"),
            on_next(130, "error"),
            on_next(220, "  foo"),
            on_next(240, " FoO "),
            on_next(270, "baR  "),
            on_next(310, "foO "),
            on_next(350, " Baz   "),
            on_next(360, "  qux "),
            on_next(390, "   bar"),
            on_next(420, " BAR  "),
            on_next(470, "FOO "),
            on_next(480, "baz  "),
            on_next(510, " bAZ "),
            on_next(530, "    fOo    "),
            on_error(570, ex),
            on_next(580, "error"),
            on_completed(600),
            on_error(650, 'ex'))
        inner_subscriptions = {}
        inners = {}
        results = {}
        c = {
            "outer_subscription": None,
            "outer": None
        }

        def action1(scheduler, state):
            c["outer"] = xs.pipe(
                    ops.group_by(
                            lambda x: x.lower().strip(),
                            lambda x: x[::-1],
                            ))

            return c["outer"]
        scheduler.schedule_absolute(created, action1)

        def action2(scheduler, state):
            def on_next(group):
                result = scheduler.create_observer()
                inners[group.key] = group
                results[group.key] = result

                def action3(scheduler, state):
                    inner_subscriptions[group.key] = group.subscribe(result, scheduler)

                scheduler.schedule_relative(100, action3)
            c["outer_subscription"] = c["outer"].subscribe(on_next, lambda e: None, scheduler=scheduler)
            return c["outer_subscription"]
        scheduler.schedule_absolute(subscribed, action2)

        def action4(scheduler, state):
            c["outer_subscription"].dispose()
            for sub in inner_subscriptions.values():
                sub.dispose()
        scheduler.schedule_absolute(disposed, action4)

        scheduler.start()
        assert len(inners) == 4
        assert results['foo'].messages == [
            on_next(470, " OOF"),
            on_next(530, "    oOf    "),
            on_error(570, ex)]
        assert results['bar'].messages == [
            on_next(390, "rab   "),
            on_next(420, "  RAB "),
            on_error(570, ex)]
        assert results['baz'].messages == [
            on_next(480, "  zab"),
            on_next(510, " ZAb "),
            on_error(570, ex)]
        assert results['qux'].messages == [
            on_error(570, ex)]
        assert xs.subscriptions == [
            subscribe(200, 570)]

    def test_group_by_with_merge(self):
        scheduler = TestScheduler()

        xs = [None]
        results = [None]

        def action1(scheduler, state):
            xs[0] = rx.from_iterable(["alpha", "apple", "beta", "bat", "gamma"]) \
                .pipe(ops.group_by(lambda s: s[0]),
                      ops.map(lambda xs: xs.pipe(ops.to_iterable(), ops.map(list))),
                      ops.merge_all(),
                      )

        scheduler.schedule_absolute(created, action1)

        def action2(scheduler, state):
            results[0] = scheduler.create_observer()
            xs[0].subscribe(results[0], scheduler)
        scheduler.schedule_absolute(subscribed, action2)

        scheduler.start()

        assert results[0].messages == [
            on_next(200, ["alpha", "apple"]),
            on_next(200, ["beta", "bat"]),
            on_next(200, ["gamma"]),
            on_completed(200)]

    def test_group_by_with_ReplaySubject(self):
        scheduler = TestScheduler()
        xs = scheduler.create_hot_observable(
            on_next(300, 1),
            on_next(310, 2),
            on_next(320, 3),
            on_next(320, 4),
            on_next(320, 5),
            on_next(320, 6),
            on_completed(1000),
            )

        observer_groups = scheduler.create_observer()
        observer_odd = scheduler.create_observer()
        observer_even = scheduler.create_observer()

        def subscription(scheduler, state):
            source = xs.pipe(
                ops.group_by(
                    key_mapper=lambda x: x % 2,
                    element_mapper=None,
                    subject_mapper=lambda: rx.subject.ReplaySubject(2),
                    )
                )
            return source.subscribe(observer_groups, scheduler=scheduler)

        scheduler.schedule_absolute(290, subscription)
        scheduler.advance_to(500)

        # extract grouped observables from messages list
        groups = {m.value.value.key: m.value.value for m in observer_groups.messages if m.value.kind =='N'}

        def subscription_odd(scheduler, state):
            source = groups[1]
            return source.subscribe(observer_odd, scheduler=scheduler)

        def subscription_even(scheduler, state):
            source = groups[0]
            return source.subscribe(observer_even, scheduler=scheduler)

        scheduler.schedule_absolute(500, subscription_odd)
        scheduler.schedule_absolute(600, subscription_even)
        scheduler.advance_to(1100)

        # only the last 2 items of odd/even are received because the
        # ReplaySubject has been configured with a buffer size of 2
        assert observer_odd.messages == [
            on_next(500, 3),
            on_next(500, 5),
            on_completed(1000),
            ]

        assert observer_even.messages == [
            on_next(600, 4),
            on_next(600, 6),
            on_completed(1000),
            ]


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