def test_complex(gen):
    constant_image = _make_constant_image()
    input = hl.ImageParam(hl.UInt(8), 3, 'input')
    input.set(constant_image)
    (x, y, c) = (hl.Var(), hl.Var(), hl.Var())
    target = hl.get_jit_target_from_environment()
    float_arg = 1.25
    int_arg = 33
    func_input = hl.Func('func_input')
    func_input[(x, y, c)] = hl.u16(((x + y) + c))
    r = gen(target, typed_buffer_input=constant_image, untyped_buffer_input=constant_image, simple_input=input, array_input=[input, input], float_arg=float_arg, int_arg=[int_arg, int_arg], untyped_buffer_output_type='uint8', extra_func_input=func_input, vectorize=True)
    (simple_output, tuple_output, array_output, typed_buffer_output, untyped_buffer_output, static_compiled_buffer_output, scalar_output, extra_func_output) = r
    b = simple_output.realize([32, 32, 3], target)
    assert (b.type() == hl.Float(32))
    for x in range(32):
        for y in range(32):
            for c in range(3):
                expected = constant_image[(x, y, c)]
                actual = b[(x, y, c)]
                assert (expected == actual), ('Expected %s Actual %s' % (expected, actual))
    b = tuple_output.realize([32, 32, 3], target)
    assert (b[0].type() == hl.Float(32))
    assert (b[1].type() == hl.Float(32))
    assert (len(b) == 2)
    for x in range(32):
        for y in range(32):
            for c in range(3):
                expected1 = (constant_image[(x, y, c)] * float_arg)
                expected2 = (expected1 + int_arg)
                (actual1, actual2) = (b[0][(x, y, c)], b[1][(x, y, c)])
                assert (expected1 == actual1), ('Expected1 %s Actual1 %s' % (expected1, actual1))
                assert (expected2 == actual2), ('Expected2 %s Actual1 %s' % (expected2, actual2))
    assert (len(array_output) == 2)
    for a in array_output:
        b = a.realize([32, 32], target)
        assert (b.type() == hl.Int(16))
        for x in range(32):
            for y in range(32):
                expected = (constant_image[(x, y, 0)] + int_arg)
                actual = b[(x, y)]
                assert (expected == actual), ('Expected %s Actual %s' % (expected, actual))
    b = typed_buffer_output.realize([32, 32, 3], target)
    assert (b.type() == hl.Float(32))
    for x in range(32):
        for y in range(32):
            for c in range(3):
                expected = constant_image[(x, y, c)]
                actual = b[(x, y, c)]
                assert (expected == actual), ('Expected %s Actual %s' % (expected, actual))
    b = untyped_buffer_output.realize([32, 32, 3], target)
    assert (b.type() == hl.UInt(8))
    for x in range(32):
        for y in range(32):
            for c in range(3):
                expected = constant_image[(x, y, c)]
                actual = b[(x, y, c)]
                assert (expected == actual), ('Expected %s Actual %s' % (expected, actual))
    b = static_compiled_buffer_output.realize([4, 4, 1], target)
    assert (b.type() == hl.UInt(8))
    for x in range(4):
        for y in range(4):
            for c in range(1):
                expected = (constant_image[(x, y, c)] + 42)
                actual = b[(x, y, c)]
                assert (expected == actual), ('Expected %s Actual %s' % (expected, actual))
    b = scalar_output.realize([], target)
    assert (b.type() == hl.Float(32))
    assert (b[()] == 34.25)
    b = extra_func_output.realize([32, 32], target)
    assert (b.type() == hl.Float(64))
    for x in range(32):
        for y in range(32):
            expected = ((x + y) + 1)
            actual = b[(x, y)]
            assert (expected == actual), ('Expected %s Actual %s' % (expected, actual))
def test_complex(gen):
    constant_image = _make_constant_image()
    input = hl.ImageParam(hl.UInt(8), 3, 'input')
    input.set(constant_image)
    (x, y, c) = (hl.Var(), hl.Var(), hl.Var())
    target = hl.get_jit_target_from_environment()
    float_arg = 1.25
    int_arg = 33
    func_input = hl.Func('func_input')
    func_input[(x, y, c)] = hl.u16(((x + y) + c))
    r = gen(target, typed_buffer_input=constant_image, untyped_buffer_input=constant_image, simple_input=input, array_input=[input, input], float_arg=float_arg, int_arg=[int_arg, int_arg], untyped_buffer_output_type='uint8', extra_func_input=func_input, vectorize=True)
    (simple_output, tuple_output, array_output, typed_buffer_output, untyped_buffer_output, static_compiled_buffer_output, scalar_output, extra_func_output) = r
    b = simple_output.realize([32, 32, 3], target)
    assert (b.type() == hl.Float(32))
    for x in range(32):
        for y in range(32):
            for c in range(3):
                expected = constant_image[(x, y, c)]
                actual = b[(x, y, c)]
                assert (expected == actual), ('Expected %s Actual %s' % (expected, actual))
    b = tuple_output.realize([32, 32, 3], target)
    assert (b[0].type() == hl.Float(32))
    assert (b[1].type() == hl.Float(32))
    assert (len(b) == 2)
    for x in range(32):
        for y in range(32):
            for c in range(3):
                expected1 = (constant_image[(x, y, c)] * float_arg)
                expected2 = (expected1 + int_arg)
                (actual1, actual2) = (b[0][(x, y, c)], b[1][(x, y, c)])
                assert (expected1 == actual1), ('Expected1 %s Actual1 %s' % (expected1, actual1))
                assert (expected2 == actual2), ('Expected2 %s Actual1 %s' % (expected2, actual2))
    assert (len(array_output) == 2)
    for a in array_output:
        b = a.realize([32, 32], target)
        assert (b.type() == hl.Int(16))
        for x in range(32):
            for y in range(32):
                expected = (constant_image[(x, y, 0)] + int_arg)
                actual = b[(x, y)]
                assert (expected == actual), ('Expected %s Actual %s' % (expected, actual))
    b = typed_buffer_output.realize([32, 32, 3], target)
    assert (b.type() == hl.Float(32))
    for x in range(32):
        for y in range(32):
            for c in range(3):
                expected = constant_image[(x, y, c)]
                actual = b[(x, y, c)]
                assert (expected == actual), ('Expected %s Actual %s' % (expected, actual))
    b = untyped_buffer_output.realize([32, 32, 3], target)
    assert (b.type() == hl.UInt(8))
    for x in range(32):
        for y in range(32):
            for c in range(3):
                expected = constant_image[(x, y, c)]
                actual = b[(x, y, c)]
                assert (expected == actual), ('Expected %s Actual %s' % (expected, actual))
    b = static_compiled_buffer_output.realize([4, 4, 1], target)
    assert (b.type() == hl.UInt(8))
    for x in range(4):
        for y in range(4):
            for c in range(1):
                expected = (constant_image[(x, y, c)] + 42)
                actual = b[(x, y, c)]
                assert (expected == actual), ('Expected %s Actual %s' % (expected, actual))
    b = scalar_output.realize([], target)
    assert (b.type() == hl.Float(32))
    assert (b[()] == 34.25)
    b = extra_func_output.realize([32, 32], target)
    assert (b.type() == hl.Float(64))
    for x in range(32):
        for y in range(32):
            expected = ((x + y) + 1)
            actual = b[(x, y)]
            assert (expected == actual), ('Expected %s Actual %s' % (expected, actual))