def test_iterator_first_maybe():
    ctx = Qit()
    fn = Function().takes(Int(), "a").returns(Int()).code("return a > 100;")
    result = ctx.run(Range(91, 120, 2).iterate().filter(fn).first_maybe())
    assert result == ("Just", 101)
    result = ctx.run(Range(10, 20, 2).iterate().filter(fn).first_maybe())
    assert result == ("Nothing", None)

def test_range_variable_generate():
    c = Qit()
    x = Variable(Int(), "x")
    r = Range(x).generate().take(30)
    result = c.run(r, { "x": 3 })
    for i in result:
        assert 0 <= i < 3

def test_iterator_first_no_default():
    ctx = Qit()
    fn = Function().takes(Int(), "a").returns(Int()).code("return a > 100;")
    result = ctx.run(Range(91, 120, 2).iterate().filter(fn).first())
    assert result == 101
    with pytest.raises(ProgramCrashed):
        ctx.run(Range(10, 20, 2).iterate().filter(fn).first())

def test_union_sort():
    ctx = Qit()
    u = Union(A=Int(), B=None, C=Int())
    v = (("A", 10), ("B", None), ("C", 5), ("B", None), ("C", 5), ("A", 0), ("A", 20))
    values = u.values(*v)
    result = ctx.run(values.iterate().sort())
    assert result == sorted(v)

def test_iterator_first_default_value():
    ctx = Qit()
    fn = Function().takes(Int(), "a").returns(Int()).code("return a > 100;")
    result = ctx.run(Range(91, 120, 2).iterate().filter(fn).first(111))
    assert result == 101
    result = ctx.run(Range(10, 20, 2).iterate().filter(fn).first(111))
    assert result == 111

def test_range_function_iterate():
    c = Qit()

    x = Variable(Int(), "x")
    r = Range(x).iterate().to_vector()
    f = r.make_function()
    assert [[], [0], [0,1], [0,1,2]] == c.run(Range(4).iterate().map(f))

def test_queue_variable():
    ctx = Qit()
    s = Struct(Int(), Int())
    q = Queue(s)
    x = Variable(q, "x")
    result = ctx.run(x, args={x: [(11,12), (5, 2)] })
    assert result == [(11, 12), (5, 2)]

def test_product_copy():
    p = Product("P", (Range(4), "x"), (Range(4), "y"))

    p2 = p.copy()
    p2.set("x", Range(2))

    q = Product("Q", (p, "p1"), (p, "p2"))

    q2 = q.copy()
    q2.set_generator("p2", p2.generator)
    q2.set_iterator("p1", p2.iterator)

    v_r4 = list(range(4))
    v_r2 = list(range(2))
    v_p = list(itertools.product(v_r4, v_r4))
    v_p2 = list(itertools.product(v_r2, v_r4))
    v_q2_generator = set(itertools.product(v_p, v_p2))
    v_q2_iterator = set(itertools.product(v_p2, v_p))

    c = Qit()
    for v in c.run(q2.generate().take(200).collect()):
        assert v in v_q2_generator

    result = c.run(q2.iterate().collect())
    assert len(v_q2_iterator) == len(result)
    assert v_q2_iterator == set(result)

def test_filter_product():
    p = Product((Range(5), "x"), (Range(5), "y"))
    q = Qit()

    f = Function("filter").returns(Bool()).takes(p.type, "p").code("return p.x == p.y;")

    q.run(p.iterate().filter(f)) == [(0, 0), (1, 1), (2, 2), (3, 3), (4, 4)]

def test_outer_variables():
    c = Qit()
    x = Variable(Int(), "x")
    f = Function().takes(Int(), "a").returns(Int()).reads(x)
    f.code("return a * x;")
    r = Range(x).iterate().map(f)
    assert [0, 5, 10, 15, 20] == c.run(r, args={ x: 5 })

def test_filter_even():
    r = Range(5)
    q = Qit()

    f = Function("filter").returns(Bool()).takes(r.type, "r").code("return r % 2 == 0;")

    assert q.run(r.iterate().filter(f)) == [0, 2, 4]

def test_vector_variable():
    ctx = Qit()
    s = Struct(Int(), Int())
    v = Vector(s)
    x = Variable(v, "x")
    result = ctx.run(x, args={x: [(11,12), (5, 2)]})
    assert result == [(11,12), (5, 2)]

def test_functor_variable():
    ctx = Qit()
    ftype = Functor("f_functor", Int(), (Int(), "x"), (Int(), "y"))

    # functions
    fplus = Function().takes(Int(), "x").takes(Int(), "y").returns(Int())
    fplus.code("return x + y;")

    ftimes = Function().takes(Int(), "x").takes(Int(), "y").returns(Int())
    ftimes.code("return x * y;")

    fmod = Function().takes(Int(), "x").takes(Int(), "y").returns(Int())
    fmod.code("return x % y;")

    # apply function in variable fvar to the given list of pairs
    fvar = ftype.variable("f")
    p = Product((Range(1, 4), "x"), (Range(1, 3), "y"))
    apply_f = Function().takes(p, "p").reads(fvar).returns(Int()).code("""
        return f(p.x, p.y);
    """)
    g = p.iterate().map(apply_f).make_function((fvar, ))

    bind_function = Function().takes(ftype, "f")\
                              .returns(Vector(ftype.return_type))
    bind_function.code("return {{g}}(f);", g=g)
    res = ctx.run(ftype.values(fplus, ftimes, fmod).iterate().map(bind_function))
    assert res == [[2, 3, 4, 3, 4, 5], [1, 2, 3, 2, 4, 6], [0, 0, 0, 1, 0, 1]]

def test_mapping_product_sequence():
    ctx = Qit()
    r = Range(2) * Range(1)
    s = Sequence(Enumerate("X", "Y"), 2)
    m = Mapping(r, s)
    result = ctx.run(m.iterate())

    expected = [
            {(1, 0): ['X', 'X'], (0, 0): ['X', 'X']},
            {(1, 0): ['X', 'X'], (0, 0): ['Y', 'X']},
            {(1, 0): ['X', 'X'], (0, 0): ['X', 'Y']},
            {(1, 0): ['X', 'X'], (0, 0): ['Y', 'Y']},
            {(1, 0): ['Y', 'X'], (0, 0): ['X', 'X']},
            {(1, 0): ['Y', 'X'], (0, 0): ['Y', 'X']},
            {(1, 0): ['Y', 'X'], (0, 0): ['X', 'Y']},
            {(1, 0): ['Y', 'X'], (0, 0): ['Y', 'Y']},
            {(1, 0): ['X', 'Y'], (0, 0): ['X', 'X']},
            {(1, 0): ['X', 'Y'], (0, 0): ['Y', 'X']},
            {(1, 0): ['X', 'Y'], (0, 0): ['X', 'Y']},
            {(1, 0): ['X', 'Y'], (0, 0): ['Y', 'Y']},
            {(1, 0): ['Y', 'Y'], (0, 0): ['X', 'X']},
            {(1, 0): ['Y', 'Y'], (0, 0): ['Y', 'X']},
            {(1, 0): ['Y', 'Y'], (0, 0): ['X', 'Y']},
            {(1, 0): ['Y', 'Y'], (0, 0): ['Y', 'Y']} ]

    def to_tuple(d):
        for key, value in d.items():
            return (tuple(key), tuple(value))

    assert set(map(to_tuple, result)) == set(map(to_tuple, expected))

def test_map_in_map():
    def prepare_map (start, size):
        return dict((start + i, start * i) for i in range(size))

    ctx = Qit()
    d = dict((i, prepare_map(i, 10)) for i in range(3))
    assert ctx.run(Map(Int(), Map(Int(), Int())).value(d)) == d

def test_struct_variable():
    ctx = Qit()
    s = Struct(Int(), Int())
    s2 = Struct(s, Int(), s)
    x = Variable(s2, "x")
    result = ctx.run(x, args={x: ((11,12), 13, (5, 2))})
    assert result == ((11,12), 13, (5, 2))

def test_system_rule_variable():
    ctx = Qit()
    x = Variable(Int(), "x")
    f = Function("f").takes(Int(), "a").returns(Int()).reads(x).code("return x;")
    s = System(Int().values(5), (f,))
    result = ctx.run(s.states(3).iterate(), args={"x": 5})
    assert result == [5]

def test_filter_empty():
    r = Range(5)
    q = Qit()

    f = Function("filter").returns(Bool()).takes(r.type, "r").code("return false;")

    assert len(q.run(r.iterate().filter(f))) == 0

def test_min_step_range_iterate_variables():
    x = Variable(Int(), "x")
    y = Variable(Int(), "y")
    z = Variable(Int(), "z")
    expr = Range(x, y, z).iterate()
    c = Qit()
    assert list(range(10, 20, 4)) == c.run(expr, args={"x": 10, "y": 20, "z": 4})

def test_vector_constructor():
    ctx = Qit()
    s = Struct(Int(), Int())
    v = Vector(s)
    x = Variable(s, "x")
    result = ctx.run(v.value([x, s.value((7, 2)), (1, 2)]),
            args={x: (11,12)})
    assert result == [ (11, 12), (7, 2), (1, 2) ]