def test_list_comprehension():
    from hidet.lang import attrs, printf

    with hidet.script_module() as script_module:

        @hidet.script
        def func():
            attrs.func_kind = 'host_kernel'
            bs = 1.0 + 1
            shape = [bs, 3, 224, 224]
            a = [1, 2, 3]
            b = [i + 1 for i in range(3) if i != 2]
            c = [s / (i + 1) for i, s in enumerate(shape)]
            printf("%d %d %d\n", a[0], a[1], a[2])
            printf("%d\n", b[0])
            printf("%f\n", c[0])
            printf("%d\n", len(b))

    func = script_module.build()
    func()

def test_dict_comprehension():
    from hidet.lang import attrs, printf

    with hidet.script_module() as script_module:

        @hidet.script
        def func():
            attrs.func_kind = 'host_kernel'
            bs = 1.0 + 1
            shape = [bs, 3, 224, 224]
            a = {k: v for k, v in enumerate(shape)}
            b = {i: (j + 1) / 3 for i, j in enumerate(range(3)) if j != 2}
            printf("%d %d %d\n", a[0], a[1], a[2])
            printf("%f\n", b[1])
            printf("%d\n", len(b))

    func = script_module.build()
    func()

def demo_softmax(shape: List[int], axis: int):
    from hidet.lang import f32
    from hidet.lang import attrs
    from hidet.lang import tensor
    from hidet.lang import grid
    import math

    with hidet.script_module() as script_module:

        @hidet.script
        def kernel(x: f32[shape], y: f32[shape]):
            attrs.func_kind = 'host_kernel'
            spatial_shape = shape[:axis] + shape[axis + 1 :]
            reduce_extent = shape[axis]

            max_value = tensor('default', f32, shape=spatial_shape)  # max(x, axis)
            exp_value = tensor('default', f32, shape=shape)  # exp(x - max)
            sum_value = tensor('default', f32, shape=spatial_shape)  # sum(exp(x - max), axis)

            # max value
            for indices in grid(spatial_shape):
                max_value[indices] = -1e10
                for k in range(reduce_extent):
                    max_value[indices] = max(max_value[indices], x[indices[:axis] + (k,) + indices[axis:]])

            # exp(x - max)
            for indices in grid(shape):
                exp_value[indices] = math.exp(x[indices] - max_value[indices[:axis] + indices[axis + 1 :]])

            # sum(exp(x - max))
            for indices in grid(spatial_shape):
                sum_value[indices] = 0.0
                for k in range(reduce_extent):
                    sum_value[indices] += exp_value[indices[:axis] + (k,) + indices[axis:]]

            # exp(x - max) / sum(exp(x - max))
            for indices in grid(shape):
                y[indices] = exp_value[indices] / sum_value[indices[:axis] + indices[axis + 1 :]]

    func = script_module.build()
    x = hidet.randn(shape)
    y1 = hidet.ops.softmax(x, axis)
    y2 = hidet.empty(shape)
    func(x, y2)
    numpy.testing.assert_allclose(y1.numpy(), y2.numpy(), rtol=1e-5, atol=1e-5)