# 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 tvm_ffi

import tvm
import tvm.script
import tvm.testing
from tvm import relax
from tvm.ir.base import assert_structural_equal
from tvm.relax import transform
from tvm.script import ir as I
from tvm.script import relax as R
from tvm.script import tirx as T


def _check_equal(x, y):
    tvm.ir.assert_structural_equal(x, y)
    tvm.ir.assert_structural_equal(y, x)

    xhash = tvm_ffi.structural_hash(x, map_free_vars=True)
    yhash = tvm_ffi.structural_hash(y, map_free_vars=True)
    assert xhash == yhash


def _check_save_roundtrip(x):
    y = tvm.ir.load_json(tvm.ir.save_json(x))
    _check_equal(x, y)


def test_basic():
    """Functions can be listed from local bindings to the IRModule"""

    # the target IRModule
    @I.ir_module(s_tir=True)
    class Expected:
        @R.function(private=True)
        def main_inner(
            x2: R.Tensor((10, 5), "float32"), y2: R.Tensor((10, 5), "float32")
        ) -> R.Tensor((10, 5), "float32"):
            s: R.Tensor((10, 5), "float32") = R.add(x2, y2)
            return s

        @R.function
        def main(x1: R.Tensor((10, 5), "float32"), y1: R.Tensor((10, 5), "float32")) -> R.Tensor(
            (10, 5), "float32"
        ):
            gv1: R.Tensor((10, 5), "float32") = Expected.main_inner(x1, y1)
            return gv1

    @I.ir_module(s_tir=True)
    class Before:
        @R.function
        def main(x1: R.Tensor((10, 5), "float32"), y1: R.Tensor((10, 5), "float32")) -> R.Tensor(
            (10, 5), "float32"
        ):
            @R.function
            def inner(
                x2: R.Tensor((10, 5), "float32"), y2: R.Tensor((10, 5), "float32")
            ) -> R.Tensor((10, 5), "float32"):
                s: R.Tensor((10, 5), "float32") = R.add(x2, y2)
                return s

            gv1: R.Tensor((10, 5), "float32") = inner(x1, y1)
            return gv1

    before = Before
    expected = Expected
    # Perform Lambda Lifting
    after = transform.LambdaLift()(before)
    assert len(after.functions) == 2
    assert_structural_equal(after, expected, map_free_vars=True)
    _check_save_roundtrip(after)


def test_input_module_is_unmodified():
    """The input module may not be modified

    If the output requires new Type, it must create a new relax
    variable.  It must not update the type of an existing relax
    variable, as that variable may be used by another IRModule.
    """

    @I.ir_module(s_tir=True)
    class Before:
        @R.function
        def main(x: R.Tensor((2, 3), "float32"), y: R.Tensor((2, 3), "float32")) -> R.Tensor(
            (2, 3), "float32"
        ):
            @R.function
            def outer_func(c1: R.Tensor((2, 3), "float32")) -> R.Callable(
                (R.Tensor((2, 3), "float32"),), R.Tensor((2, 3), "float32")
            ):
                @R.function
                def inner_func(x1: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "float32"):
                    s: R.Tensor((2, 3), "float32") = R.add(x1, c1)
                    return s

                return inner_func

            in_call = outer_func(x)
            res = in_call(y)
            return res

    before = Before
    copy_of_before = tvm.ir.load_json(tvm.ir.save_json(before))

    transform.LambdaLift()(before)

    tvm.ir.assert_structural_equal(before, copy_of_before)


def test_closure():
    """Lifting functions may require producing closures"""

    # the expected IRModule
    @I.ir_module(s_tir=True)
    class Expected:
        @R.function
        def main(x: R.Tensor((2, 3), "float32"), y: R.Tensor((2, 3), "float32")) -> R.Tensor(
            (2, 3), "float32"
        ):
            in_call = Expected.main_outer_func(x)
            res = R.invoke_pure_closure(in_call, (y,), ty_args=(R.Tensor((2, 3), dtype="float32")))
            return res

        @R.function(private=True)
        def main_inner_func(x1: R.Tensor((2, 3), "float32"), c1: R.Tensor((2, 3), "float32")):
            r_1: R.Tensor((2, 3), "float32") = R.add(x1, c1)
            return r_1

        @R.function(private=True)
        def main_outer_func(y: R.Tensor((2, 3), "float32")) -> R.Any:
            inner_func = R.make_closure(Expected.main_inner_func, (y,))
            return inner_func

    # IRModule to perform Lambda Lifting
    @I.ir_module(s_tir=True)
    class Before:
        @R.function
        def main(x: R.Tensor((2, 3), "float32"), y: R.Tensor((2, 3), "float32")) -> R.Tensor(
            (2, 3), "float32"
        ):
            @R.function
            def outer_func(c1: R.Tensor((2, 3), "float32")) -> R.Callable(
                (R.Tensor((2, 3), "float32"),), R.Tensor((2, 3), "float32")
            ):
                @R.function
                def inner_func(x1: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "float32"):
                    s: R.Tensor((2, 3), "float32") = R.add(x1, c1)
                    return s

                return inner_func

            in_call = outer_func(x)
            res = in_call(y)
            return res

    before = Before
    after = transform.LambdaLift()(before)
    expected = Expected
    assert_structural_equal(after, expected, map_free_vars=True)
    _check_save_roundtrip(after)


def test_recursive():
    """The lifted function may be recursively defined"""

    # the expected IRModule
    @I.ir_module(s_tir=True)
    class Expected:
        @R.function(private=True)
        def main_while_loop(
            i: R.Tensor((), "int32"), s: R.Tensor((2, 3), "float32"), x: R.Tensor((2, 3), "float32")
        ) -> R.Tensor((2, 3), "float32"):
            cond: R.Tensor((), "bool") = R.call_pure_packed(
                "test.vm.less", i, R.const(10), ty_args=(R.Tensor((), dtype="bool"))
            )
            c: R.Tensor((), "int32") = R.const(1, dtype="int32")
            if cond:
                new_i: R.Tensor((), "int32") = R.add(i, c)
                new_s: R.Tensor((2, 3), "float32") = R.add(s, x)
                new_r = Expected.main_while_loop(new_i, new_s, x)
                r = new_r
            else:
                r = s
            return r

        @R.function
        def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), dtype="float32"):
            while_loop = R.make_closure(Expected.main_while_loop, (x,))
            gv: R.Tensor((2, 3), dtype="float32") = R.invoke_pure_closure(
                while_loop,
                (R.const(0), x),
                ty_args=(R.Tensor((2, 3), dtype="float32")),
            )
            return gv

    # the IRModule to apply lambda lifting
    @I.ir_module(s_tir=True)
    class Before:
        @R.function
        def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor:
            @R.function
            def while_loop(i: R.Tensor((), "int32"), s: R.Tensor((2, 3), "float32")) -> R.Tensor(
                (2, 3), "float32"
            ):
                cond: R.Tensor((), "bool") = R.call_pure_packed(
                    "test.vm.less", i, R.const(10), ty_args=(R.Tensor((), dtype="bool"))
                )
                c: R.Tensor((), "int32") = R.const(1, dtype="int32")
                if cond:
                    new_i: R.Tensor((), "int32") = R.add(i, c)
                    new_s: R.Tensor((2, 3), "float32") = R.add(s, x)
                    r: R.Tensor((2, 3), "float32") = while_loop(new_i, new_s)
                else:
                    r: R.Tensor((2, 3), "float32") = s
                return r

            gv: R.Tensor((2, 3), "float32") = while_loop(R.const(0), x)
            return gv

    before = Before
    expected = Expected
    # check well-formness of recursive call
    relax.analysis.well_formed(before)

    # Perform Lambda Lifting
    after = transform.LambdaLift()(before)
    assert len(after.functions) == 2

    assert_structural_equal(after, expected, map_free_vars=True)
    _check_save_roundtrip(after)


def test_multi_func():
    """Lifting may be required for multiple top-level functions

    De-duplication of GlobalVar names at the IRModule is done by
    appending the name of the function from which they were lifted.
    """

    # expected IRModule
    @I.ir_module(s_tir=True)
    class Expected:
        @R.function
        def glob_func_1(
            x1: R.Tensor((10, 5), "float32"), y1: R.Tensor((10, 5), "float32")
        ) -> R.Tensor(None, "float32", ndim=2):
            gv1: R.Tensor((10, 5), "float32") = Expected.glob_func_1_inner(x1, y1)
            return gv1

        @R.function
        def glob_func_2(
            x11: R.Tensor((10, 5), "float32"), y11: R.Tensor((10, 5), "float32")
        ) -> R.Tensor(None, "float32", ndim=2):
            gv11: R.Tensor((10, 5), "float32") = Expected.glob_func_2_inner(x11, y11)
            return gv11

        @R.function(private=True)
        def glob_func_1_inner(
            x2: R.Tensor((10, 5), "float32"), y2: R.Tensor((10, 5), "float32")
        ) -> R.Tensor((10, 5), "float32"):
            s: R.Tensor((10, 5), "float32") = R.add(x2, y2)
            return s

        @R.function(private=True)
        def glob_func_2_inner(
            x21: R.Tensor((10, 5), "float32"), y21: R.Tensor((10, 5), "float32")
        ) -> R.Tensor((10, 5), "float32"):
            s1: R.Tensor((10, 5), "float32") = R.add(x21, y21)
            return s1

    # the IRModule to apply lambda lifting
    @I.ir_module(s_tir=True)
    class Before:
        @R.function
        def glob_func_1(
            x1: R.Tensor((10, 5), "float32"), y1: R.Tensor((10, 5), "float32")
        ) -> R.Tensor((10, 5), "float32"):
            @R.function
            def inner(
                x2: R.Tensor((10, 5), "float32"), y2: R.Tensor((10, 5), "float32")
            ) -> R.Tensor((10, 5), "float32"):
                s: R.Tensor((10, 5), "float32") = R.add(x2, y2)
                return s

            gv1: R.Tensor((10, 5), "float32") = inner(x1, y1)
            return gv1

        @R.function
        def glob_func_2(
            x1: R.Tensor((10, 5), "float32"), y1: R.Tensor((10, 5), "float32")
        ) -> R.Tensor((10, 5), "float32"):
            @R.function
            def inner(
                x2: R.Tensor((10, 5), "float32"), y2: R.Tensor((10, 5), "float32")
            ) -> R.Tensor((10, 5), "float32"):
                s: R.Tensor((10, 5), "float32") = R.add(x2, y2)
                return s

            gv1: R.Tensor((10, 5), "float32") = inner(x1, y1)
            return gv1

    before = Before
    expected = Expected
    # Perform Lambda Lifting
    after = transform.LambdaLift()(before)
    assert len(after.functions) == 4
    assert_structural_equal(after, expected, map_free_vars=True)
    _check_save_roundtrip(after)


def test_no_local_func():
    @I.ir_module(s_tir=True)
    class Before:
        @T.prim_func(s_tir=True)
        def sub(
            A: T.Buffer((16, 16), "float32"),
            B: T.Buffer((16, 16), "float32"),
            C: T.Buffer((16, 16), "float32"),
        ) -> None:
            for i, j in T.grid(16, 16):
                with T.sblock("sub"):
                    vi, vj = T.axis.remap("SS", [i, j])
                    C[vi, vj] = A[vi, vj] - B[vi, vj]

        @R.function
        def before(c0: R.Tensor((16, 16), "float32"), x: R.Tensor(dtype="float32", ndim=2)):
            s = R.call_tir(Before.sub, (c0, x), R.Tensor((16, 16), dtype="float32"))
            return s

    before = Before
    # Perform lambda lifting
    after = transform.LambdaLift()(before)
    # No local functions are lifted
    assert_structural_equal(after, before, map_free_vars=True)
    _check_save_roundtrip(after)


def test_impure_function():
    @I.ir_module(s_tir=True)
    class Expected:
        @R.function(pure=False, private=True)
        def main_inner() -> R.Tuple:
            y = R.print(format="Wow!")
            return y

        @R.function(pure=False)
        def main(x: R.Tensor((), "int32")) -> R.Tensor((), "int32"):
            gv1 = Expected.main_inner()
            return x

    @I.ir_module(s_tir=True)
    class Before:
        @R.function(pure=False)
        def main(x: R.Tensor((), "int32")) -> R.Tensor((), "int32"):
            @R.function(pure=False)
            def inner() -> R.Tuple:
                y = R.print(format="Wow!")
                return y

            gv1 = inner()
            return x

    before = Before
    expected = Expected
    after = transform.LambdaLift()(before)
    assert len(after.functions) == 2
    assert_structural_equal(after, expected, map_free_vars=True)
    _check_save_roundtrip(after)


def test_lambda_function_with_same_name_as_global():
    """Lifted lambda names may not conflict with previous names

    Like `test_basic`, but the module has an existing function
    `main_inner`, which has the same name as the LambdaLift's first
    choice of name for the hoisted function.
    """

    @I.ir_module(s_tir=True)
    class Before:
        @R.function
        def main(x1: R.Tensor((10, 5), "float32"), y1: R.Tensor((10, 5), "float32")) -> R.Tensor(
            (10, 5), "float32"
        ):
            @R.function
            def inner(
                x2: R.Tensor((10, 5), "float32"), y2: R.Tensor((10, 5), "float32")
            ) -> R.Tensor((10, 5), "float32"):
                s: R.Tensor((10, 5), "float32") = R.add(x2, y2)
                return s

            gv1: R.Tensor((10, 5), "float32") = inner(x1, y1)
            return gv1

        @R.function
        def main_inner():
            return R.tuple()

    @I.ir_module(s_tir=True)
    class Expected:
        @R.function
        def main(x1: R.Tensor((10, 5), "float32"), y1: R.Tensor((10, 5), "float32")) -> R.Tensor(
            (10, 5), "float32"
        ):
            gv1: R.Tensor((10, 5), "float32") = Expected.main_inner_0(x1, y1)
            return gv1

        @R.function(private=True)
        def main_inner_0(
            x2: R.Tensor((10, 5), "float32"), y2: R.Tensor((10, 5), "float32")
        ) -> R.Tensor((10, 5), "float32"):
            s: R.Tensor((10, 5), "float32") = R.add(x2, y2)
            return s

        @R.function
        def main_inner():
            return R.tuple()

    after = transform.LambdaLift()(Before)
    assert_structural_equal(Expected, after)


def test_symbolic_variable_defined_by_inner_func():
    @I.ir_module(s_tir=True)
    class Before:
        @R.function
        def main(x1: R.Tensor((10, 5), "float32"), y1: R.Tensor((10, 5), "float32")) -> R.Tensor(
            (10, 5), "float32"
        ):
            @R.function
            def inner(x2: R.Tensor(("n", "m"), "float32"), y2: R.Tensor(("n", "m"), "float32")):
                sum_inner = R.add(x2, y2)
                return sum_inner

            sum_main = inner(x1, y1)
            return sum_main

    @I.ir_module(s_tir=True)
    class Expected:
        @R.function
        def main(x1: R.Tensor((10, 5), "float32"), y1: R.Tensor((10, 5), "float32")) -> R.Tensor(
            (10, 5), "float32"
        ):
            sum_main = Expected.main_inner(x1, y1)
            return sum_main

        @R.function(private=True)
        def main_inner(
            x2: R.Tensor(("n", "m"), "float32"), y2: R.Tensor(("n", "m"), "float32")
        ) -> R.Tensor(("n", "m"), "float32"):
            sum_inner = R.add(x2, y2)
            return sum_inner

    After = transform.LambdaLift()(Before)
    assert_structural_equal(Expected, After)


def test_symbolic_variable_defined_by_outer_func():
    @I.ir_module(s_tir=True)
    class Before:
        @R.function
        def main(
            x1: R.Tensor(("n", "m"), "float32"), y1: R.Tensor(("n", "m"), "float32")
        ) -> R.Tensor(("n", "m"), "float32"):
            n = T.int64()
            m = T.int64()

            @R.function
            def inner(x2: R.Tensor((n, m), "float32"), y2: R.Tensor((n, m), "float32")):
                sum_inner = R.add(x2, y2)
                return sum_inner

            sum_main = inner(x1, y1)
            return sum_main

    @I.ir_module(s_tir=True)
    class Expected:
        @R.function
        def main(
            x1: R.Tensor(("n", "m"), "float32"), y1: R.Tensor(("n", "m"), "float32")
        ) -> R.Tensor(("n", "m"), "float32"):
            sum_main = Expected.main_inner(x1, y1)
            return sum_main

        @R.function(private=True)
        def main_inner(
            x2: R.Tensor(("n", "m"), "float32"), y2: R.Tensor(("n", "m"), "float32")
        ) -> R.Tensor(("n", "m"), "float32"):
            sum_inner = R.add(x2, y2)
            return sum_inner

    After = transform.LambdaLift()(Before)
    assert_structural_equal(Expected, After)


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