# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
# ruff: noqa: F841
import pytest

import tvm
import tvm.script
import tvm.testing
from tvm import relax, tir
from tvm.ir.base import assert_structural_equal
from tvm.script import relax as R
from tvm.script import tir as T


def test_normalize_function():
    m = tir.Var("m", "int64")
    n = tir.Var("n", "int64")
    x = relax.Var("x", R.Tensor([m, n], "float16"))

    # Note: the parser automatically normalize the IR written in TVMScript,
    # so we manually construct the function here.
    mul_add = relax.Function(
        [x],
        relax.op.multiply(relax.op.add(x, x), relax.op.add(x, x)),
        ret_struct_info=R.Tensor("float16", ndim=2),
    )

    # Note: from_expr api names private function (function without global_symbol) as "main"
    before_mod = tvm.IRModule.from_expr(mul_add)

    after_mod = relax.transform.Normalize()(before_mod)

    @R.function(private=True)
    def expected(x: R.Tensor(("m", "n"), "float16")) -> R.Tensor(dtype="float16", ndim=2):
        gv = R.add(x, x)
        gv1 = R.add(x, x)
        return R.multiply(gv, gv1)

    assert_structural_equal(after_mod["main"], expected)


def test_normalize_if():
    cond = relax.Var("cond", R.Tensor([], "bool"))
    x = relax.Var("x", R.Tensor([1], "float32"))
    # TODO(relax-team): add type and shape inference for IfNode
    y = relax.Var("y")

    # Note: the parser automatically normalize the IR written in TVMScript,
    # so we manually construct the function and If here.
    f = relax.Function(
        [cond, x],
        relax.SeqExpr(
            [
                relax.BindingBlock(
                    [
                        relax.VarBinding(
                            y,
                            relax.If(
                                cond,
                                relax.op.multiply(relax.op.add(x, x), relax.op.add(x, x)),
                                relax.op.add(relax.op.multiply(x, x), relax.op.multiply(x, x)),
                            ),
                        )
                    ]
                )
            ],
            y,
        ),
        ret_struct_info=R.Tensor("float32", ndim=1),
    )

    before_mod = tvm.IRModule.from_expr(f)
    after_mod = relax.transform.Normalize()(before_mod)

    @R.function(private=True)
    def expected(cond: R.Tensor((), "bool"), x: R.Tensor((1,), "float32")) -> R.Tensor(
        dtype="float32", ndim=1
    ):
        if cond:
            gv = R.add(x, x)
            gv1 = R.add(x, x)
            y = R.multiply(gv, gv1)
        else:
            gv = R.multiply(x, x)
            gv1 = R.multiply(x, x)
            y = R.add(gv, gv1)
        return y

    assert_structural_equal(after_mod["main"], expected)


def test_normalize_no_op():
    # the normalize pass should be no-op for IR in ANF
    @tvm.script.ir_module
    class ANFMod1:
        @R.function
        def f(x: R.Tensor(dtype="float32")):
            gv = R.add(x, x)
            gv1 = R.add(gv, gv)
            gv2 = R.add(gv, gv1)
            return (gv, gv2)

    before_mod = ANFMod1
    after_mod = relax.transform.Normalize()(before_mod)
    assert_structural_equal(before_mod, after_mod, map_free_vars=True)

    @tvm.script.ir_module
    class ANFMod2:
        @R.function
        def foo(x: R.Tensor(("m", "n"), "float32")):
            m, n = T.int64(), T.int64()
            with R.dataflow():
                lv0 = R.call_dps_packed("test.op.identity", (x,), R.Tensor((m, n), dtype="float32"))
                gv0 = R.call_dps_packed(
                    "test.op.identity", (lv0,), R.Tensor((m, n), dtype="float32")
                )
                R.output(gv0)
            return gv0

    mod = ANFMod2
    mod_post = relax.transform.Normalize()(mod)

    assert_structural_equal(mod, mod_post)


def test_normalize_seq_body():
    # a seq expression with a non-leaf body should bind the body to a var as well
    x = relax.Var("x", R.Tensor([], "int32"))
    y = relax.Var("y", R.Tensor([], "int32"))
    seq = relax.SeqExpr([], relax.op.add(x, y))
    f = relax.Function(
        [x, y],
        seq,
        ret_struct_info=R.Tensor([], "int32"),
    )

    before_mod = tvm.IRModule.from_expr(f)
    after_mod = relax.transform.Normalize()(before_mod)

    @R.function(private=True)
    def expected(x: R.Tensor((), dtype="int32"), y: R.Tensor((), dtype="int32")) -> R.Tensor(
        ndim=0, dtype="int32"
    ):
        # normalization inserts a binding like this
        z = R.add(x, y)
        return z

    assert_structural_equal(after_mod["main"], expected)


def test_normalize_func_body():
    # a function with a body that is not a seq expr should have it wrapped in a seq expr
    x = relax.Var("x", R.Tensor([], "int32"))
    y = relax.Var("y", R.Tensor([], "int32"))
    f = relax.Function(
        [x, y],
        relax.op.add(x, y),
        ret_struct_info=R.Tensor([], "int32"),
    )

    before_mod = tvm.IRModule.from_expr(f)
    after_mod = relax.transform.Normalize()(before_mod)

    @R.function(private=True)
    def expected(x: R.Tensor((), dtype="int32"), y: R.Tensor((), dtype="int32")) -> R.Tensor(
        ndim=0, dtype="int32"
    ):
        # result will be a seq expr where the body is a var
        z = R.add(x, y)
        return z

    assert_structural_equal(after_mod["main"], expected)


def test_normalize_if_branches():
    # an if node's branches must be seq exprs
    x = relax.Var("x", R.Tensor([], "int32"))
    y = relax.Var("y", R.Tensor([], "int32"))
    # TODO(@relax-team): z has a shape of () and type of TensorType(ndim=0),
    # but normalization fails to infer these even though it should
    z = relax.Var("z")
    cond = relax.Var("cond", R.Tensor([], "bool"))
    plus = relax.op.add(x, y)
    mult = relax.op.multiply(x, y)
    if_node = relax.If(cond, plus, mult)
    seq = relax.SeqExpr([relax.BindingBlock([relax.VarBinding(z, if_node)])], z)
    f = relax.Function(
        [cond, x, y],
        seq,
        ret_struct_info=R.Tensor([], "int32"),
    )

    before_mod = tvm.IRModule.from_expr(f)
    after_mod = relax.transform.Normalize()(before_mod)

    @R.function(private=True)
    def expected(
        cond: R.Tensor((), dtype="bool"),
        x: R.Tensor((), dtype="int32"),
        y: R.Tensor((), dtype="int32"),
    ) -> R.Tensor(ndim=0, dtype="int32"):
        # the bodies of the branches will be seq exprs with a binding
        if cond:
            w = R.add(x, y)
            z = w
        else:
            w = R.multiply(x, y)
            z = w
        return z

    assert_structural_equal(after_mod["main"], expected)


def test_normalize_if_condition():
    cond = relax.Var("cond", R.Tensor([], "bool"))
    x = relax.Var("x", R.Tensor([1], "float32"))
    # TODO(relax-team): add type and shape inference for IfNode
    y = relax.Var("y")

    # The condition is wrapped in a tuple and then indexed
    f = relax.Function(
        [cond, x],
        relax.SeqExpr(
            [
                relax.BindingBlock(
                    [
                        relax.VarBinding(
                            y,
                            relax.If(
                                relax.TupleGetItem(relax.Tuple([cond]), 0),
                                relax.op.add(x, x),
                                relax.op.multiply(x, x),
                            ),
                        )
                    ]
                )
            ],
            y,
        ),
        ret_struct_info=R.Tensor("float32", ndim=1),
    )

    before_mod = tvm.IRModule.from_expr(f)
    after_mod = relax.transform.Normalize()(before_mod)

    @R.function(private=True)
    def expected(cond: R.Tensor((), "bool"), x: R.Tensor((1,), "float32")) -> R.Tensor(
        dtype="float32", ndim=1
    ):
        c = R.TupleGetItem(R.tuple(cond), 0)
        if c:
            gv = R.add(x, x)
            y = gv
        else:
            gv = R.multiply(x, x)
            y = gv
        return y

    assert_structural_equal(after_mod["main"], expected)


def test_normalize_tuple_get_item():
    x = relax.Var("x", R.Tensor([], "int32"))
    f = relax.Function(
        [x],
        relax.TupleGetItem(
            relax.TupleGetItem(
                relax.Tuple([relax.Tuple([x])]),
                0,
            ),
            0,
        ),
        ret_struct_info=R.Tensor([], "int32"),
    )

    before_mod = tvm.IRModule.from_expr(f)
    after_mod = relax.transform.Normalize()(before_mod)

    # TODO: Revisit once we canonicalize SeqExprs (part of normalization?)
    # Not using the parser this time because writing it out correctly results in
    # *one* binding block, whereas the normalized version has *two*
    idx_var = relax.Var("idx_var", R.Tuple([R.Tensor([], "int32")]))
    ret_var = relax.Var("ret", R.Tensor([], "int32"))
    expected_f = relax.Function(
        [x],
        relax.SeqExpr(
            [
                relax.BindingBlock(
                    [
                        relax.VarBinding(
                            idx_var, relax.TupleGetItem(relax.Tuple([relax.Tuple([x])]), 0)
                        )
                    ]
                ),
                relax.BindingBlock([relax.VarBinding(ret_var, relax.TupleGetItem(idx_var, 0))]),
            ],
            ret_var,
        ),
        ret_struct_info=R.Tensor([], "int32"),
    )
    expected_mod = tvm.IRModule.from_expr(expected_f)
    # apply normalization to fill in type and shape annotations (tedious otherwise)
    final_mod = relax.transform.Normalize()(expected_mod)

    assert_structural_equal(after_mod, final_mod)


def test_normalize_combine_nearby_blocks():
    x = relax.Var("x", R.Tensor([], "int32"))
    v0 = relax.Var("v0", R.Tensor([], "int32"))
    v1 = relax.Var("v1", R.Tensor([], "int32"))
    v2 = relax.Var("v2", R.Tensor([], "int32"))
    v3 = relax.Var("v3", R.Tensor([], "int32"))
    f = relax.Function(
        [x],
        relax.SeqExpr(
            [
                relax.DataflowBlock([relax.VarBinding(v0, x)]),
                relax.DataflowBlock([relax.VarBinding(v1, v0)]),
                relax.BindingBlock([relax.VarBinding(v2, v1)]),
                relax.BindingBlock([relax.VarBinding(v3, v2)]),
            ],
            v3,
        ),
        ret_struct_info=R.Tensor([], "int32"),
    )

    after_mod = relax.transform.Normalize()(tvm.IRModule.from_expr(f))

    @R.function(private=True)
    def expected(x: R.Tensor((), "int32")):
        with R.dataflow():
            v0 = x
            v1 = v0
            R.output(v0, v1)
        v2 = v1
        v3 = v2
        return v3

    assert_structural_equal(after_mod["main"], expected)


def test_normalize_nested_seq():
    x = relax.Var("x", R.Tensor([], "int32"))
    y = relax.Var("y", R.Tensor([], "int32"))
    z = relax.Var("z", R.Tensor([], "int32"))
    seq = relax.SeqExpr(
        [
            relax.BindingBlock(
                [
                    relax.VarBinding(x, relax.const(1)),
                    relax.VarBinding(
                        y,
                        relax.SeqExpr(
                            [relax.BindingBlock([relax.VarBinding(z, relax.const(2))])],
                            z,
                        ),
                    ),
                ]
            )
        ],
        y,
    )

    f = relax.Function(
        [],
        seq,
        ret_struct_info=R.Tensor([], "int32"),
    )
    after_mod = relax.transform.Normalize()(tvm.IRModule.from_expr(f))

    @R.function(private=True)
    def expected():
        x = relax.const(1)
        z = relax.const(2)
        y = z
        return y

    assert_structural_equal(after_mod["main"], expected)


def test_normalize_nested_seq_dataflow():
    x = relax.Var("x", R.Tensor([], "int32"))
    y = relax.Var("y", R.Tensor([], "int32"))
    z = relax.Var("z", R.Tensor([], "int32"))
    q = relax.Var("u", R.Tensor([], "int32"))
    w = relax.DataflowVar("w", R.Tensor([], "int32"))
    u = relax.Var("u", R.Tensor([], "int32"))
    seq = relax.SeqExpr(
        [
            relax.BindingBlock(
                [
                    relax.VarBinding(x, relax.const(1)),
                    relax.VarBinding(
                        y,
                        relax.SeqExpr(
                            [
                                relax.BindingBlock([relax.VarBinding(q, relax.const(2))]),
                                relax.DataflowBlock(
                                    [
                                        relax.VarBinding(w, q),
                                        relax.VarBinding(u, w),
                                    ]
                                ),
                                relax.BindingBlock([relax.VarBinding(z, u)]),
                            ],
                            z,
                        ),
                    ),
                ]
            )
        ],
        y,
    )

    f = relax.Function(
        [],
        seq,
        ret_struct_info=R.Tensor([], "int32"),
    )
    after_mod = relax.transform.Normalize()(tvm.IRModule.from_expr(f))

    @R.function(private=True)
    def expected():
        x = relax.const(1)
        q = relax.const(2)
        with R.dataflow():
            w = q
            u = w
            R.output(u)
        z = u
        y = z
        return y

    assert_structural_equal(after_mod["main"], expected)


def test_normalize_deeply_nested_seq():
    x = relax.Var("x", R.Tensor([], "int32"))
    y = relax.Var("y", R.Tensor([], "int32"))
    z = relax.Var("z", R.Tensor([], "int32"))
    u = relax.Var("u", R.Tensor([], "int32"))
    v = relax.Var("v", R.Tensor([], "int32"))
    w = relax.Var("w", R.Tensor([], "int32"))
    _ = relax.Var("w", R.Tensor([], "int32"))
    seq = relax.SeqExpr(
        [
            relax.BindingBlock(
                [
                    relax.VarBinding(x, relax.const(1)),
                    relax.VarBinding(
                        y,
                        relax.SeqExpr(
                            [
                                relax.BindingBlock(
                                    [
                                        relax.VarBinding(
                                            z,
                                            relax.SeqExpr(
                                                [
                                                    relax.BindingBlock(
                                                        [
                                                            relax.VarBinding(u, relax.const(2)),
                                                            relax.MatchCast(
                                                                _, u, R.Tensor([], "int32")
                                                            ),
                                                            relax.VarBinding(v, u),
                                                            relax.MatchCast(
                                                                w, v, R.Tensor([], "int32")
                                                            ),
                                                        ]
                                                    )
                                                ],
                                                w,
                                            ),
                                        )
                                    ]
                                )
                            ],
                            z,
                        ),
                    ),
                ]
            )
        ],
        y,
    )

    f = relax.Function(
        [],
        seq,
        ret_struct_info=R.Tensor([], "int32"),
    )
    after_mod = relax.transform.Normalize()(tvm.IRModule.from_expr(f))

    @R.function(private=True)
    def expected():
        x = relax.const(1)
        u = relax.const(2)
        _ = R.match_cast(u, R.Tensor((), "int32"))
        v = u
        w = R.match_cast(v, R.Tensor((), "int32"))
        z = w
        y = z
        return y

    assert_structural_equal(after_mod["main"], expected)


@pytest.mark.xfail()
def test_nesting_non_dataflow_in_dataflow_error():
    x = relax.DataflowVar("x", R.Tensor([], "int32"))
    y = relax.Var("y", R.Tensor([], "int32"))
    z = relax.Var("z", R.Tensor([], "int32"))
    seq = relax.SeqExpr(
        [
            relax.DataflowBlock(
                [
                    relax.VarBinding(x, relax.const(1)),
                    relax.VarBinding(
                        y,
                        relax.SeqExpr(
                            [relax.BindingBlock([relax.VarBinding(z, relax.const(2))])],
                            z,
                        ),
                    ),
                ]
            )
        ],
        y,
    )
    f = relax.Function(
        [],
        seq,
        ret_struct_info=R.Tensor([], "int32"),
    )
    relax.transform.Normalize()(tvm.IRModule.from_expr(f))
    # should fail due to a normal binding block being inside a dataflowblock


def test_remove_usage_of_void_type_variables():
    """All empty tuples should be constructed in-line

    For readability, TVMScript hides the variable binding if the
    variable has a void type.  For example, `R.assert_op(condition)`
    instead of `void_var: R.Tuple([]) = R.assert_op(condition)`.
    However, Relax follows standard convention of functional
    languages, and uses an empty tuple to represent void.  Since an
    empty tuple may be legally used later in the function, the
    `void_var` may require a binding.

    This is avoided by normalizing all usage of a void-type
    variable with an in-line `R.tuple()`.
    """
    x = relax.Var("x", R.Tuple([]))
    bindings = [
        relax.VarBinding(x, R.assert_op(R.const(True, "bool"))),
    ]
    seq = relax.SeqExpr([relax.BindingBlock(bindings)], x)
    before = relax.Function([], seq, ret_struct_info=R.Tuple([]), is_pure=False)

    after = relax.transform.Normalize()(tvm.IRModule({"main": before}))["main"]

    @R.function(private=True, pure=False)
    def expected():
        x = R.assert_op(R.const(True, "bool"))
        return R.tuple()


if __name__ == "__main__":
    tvm.testing.main()