# 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.
# pylint: disable=invalid-name, unused-argument, missing-function-docstring, abstract-method, missing-class-docstring
# ruff: noqa: RUF005
"""Relax LazyTransformParams pass."""

import tvm
from tvm import IRModule, relax
from tvm.relax.expr_functor import PyExprMutator, PyExprVisitor, mutator, visitor


@visitor
class ForwardCollector(PyExprVisitor):
    """
    Perform a forward pass to collect the following information:
    out_tuple_map: map from var to its index in the output tuple
    var_tuple_get_item: list of var that is bound to v = params[i]

    Parameters
    ----------
    tuple_var: relax.Var
        The output tuple var

    input_params: relax.Var
        The input tuple var

    """

    def __init__(self, tuple_var: relax.Var, input_params: relax.Var) -> None:
        self.out_tuple_map = {}
        self.out_tuple_var = tuple_var
        self.input_params = input_params
        self.var_tuple_get_item = []
        self.is_tuple_get_item_input = False

    def visit_tuple_getitem_(self, op: relax.TupleGetItem) -> None:
        if op.tuple_value == self.input_params:
            self.is_tuple_get_item_input = True
        else:
            self.is_tuple_get_item_input = False
        super().visit_tuple_getitem_(op)

    def visit_var_binding_(self, binding: relax.VarBinding) -> None:
        if binding.var == self.out_tuple_var:
            assert isinstance(binding.value, relax.Tuple)
            for i, expr in enumerate(binding.value.fields):
                if expr not in self.out_tuple_map:
                    self.out_tuple_map[expr] = []
                self.out_tuple_map[expr].append(relax.prim_value(i))
        else:
            self.is_tuple_get_item_input = False
            super().visit_var_binding_(binding)
            if self.is_tuple_get_item_input:
                self.var_tuple_get_item.append(binding.var)


@visitor
class LivenessAnalysis(PyExprVisitor):
    """
    Perform a backward pass to collect the following information:
    var_liveness_end: map from var to the list of var whose liveness is killed by this var binding

    Parameters
    ----------
    out_tuple_var: relax.Var
        The output tuple var
    input_params: set
        The set of vars that are bound to v = params[i]
    """

    def __init__(self, out_tuple_var: relax.Var) -> None:
        self.last_appear_in_var_binding = []
        self.out_tuple_var = out_tuple_var
        self.var_liveness_end = {}
        self.ended_vars = set()

    def visit_binding_block_(self, block: relax.BindingBlock) -> None:
        for binding in reversed(block.bindings):
            self.visit_binding(binding)

    def visit_var_(self, op: relax.Var) -> None:
        if op not in self.ended_vars:
            self.last_appear_in_var_binding.append(op)
            self.ended_vars.add(op)

    def visit_var_binding_(self, binding: relax.VarBinding) -> None:
        if self.out_tuple_var == binding.var:
            return
        self.last_appear_in_var_binding = []
        super().visit_var_binding_(binding)
        # param[i] is in output
        if binding.var not in self.ended_vars:
            self.last_appear_in_var_binding.append(binding.var)
            self.ended_vars.add(binding.var)
        self.var_liveness_end[binding.var] = self.last_appear_in_var_binding


class LazyTransformParamsFuncCreator:
    """
    Transform transform_params functions into a lazy version.

    Parameters
    ----------
    mod: IRModule
        The module to be transformed
    """

    def __init__(
        self,
        fget_item,
        fset_item,
        extra_get_item_params,
        extra_set_item_params,
        mod: IRModule = None,
    ) -> None:
        self.mod = mod
        self.fget_item = fget_item
        self.extra_get_item_params = extra_get_item_params
        self.fset_item = fset_item
        self.extra_set_item_params = extra_set_item_params
        self.input_params_set = None
        self.out_tuple_map = None
        self.out_tuple_var = None
        self.memory_free_insertion = None

    def transform(self, func: relax.Function) -> relax.Function:
        if "num_input" in func.attrs:
            num_input = func.attrs["num_input"]
        else:
            num_input = 0

        seq_expr = func.body
        self.out_tuple_var = seq_expr.body

        # Step 1. collect out_tuple_map and input_params_set
        forward_collector = ForwardCollector(self.out_tuple_var, func.params[num_input])
        forward_collector.visit_expr(func)
        self.out_tuple_map = forward_collector.out_tuple_map
        # input_params_set is the set of binding var for var = params[i]
        self.input_params_set = set(forward_collector.var_tuple_get_item)

        # Step 2. liveness analysis and get where to insert kill_object instruction
        liveness = LivenessAnalysis(self.out_tuple_var)
        liveness.visit_expr(func)
        self.memory_free_insertion = liveness.var_liveness_end

        # Step 3. rewrite get item and set item
        if self.fget_item is not None:
            new_func = LazyInputMutator(self, self.mod).visit_expr(func)

        new_body = new_func.body
        if self.fset_item is not None:
            # The LazyOutputMutator only inspects variable bindings
            # for replacement.  If the output tuple includes elements
            # that do not have a variable binding, such as
            # `relax.Const`, these must still produce a call to the
            # `"set_item"` function.
            leaf_outputs = {
                expr: indices
                for expr, indices in self.out_tuple_map.items()
                if not isinstance(expr, relax.Var)
            }
            if leaf_outputs:
                new_bindings = [
                    relax.VarBinding(
                        relax.Var("_", relax.AnyType()),
                        relax.Call(
                            relax.ExternFunc(self.fset_item),
                            [*self.extra_set_item_params, index, expr],
                            None,
                            [relax.AnyType()],
                        ),
                    )
                    for expr, indices in leaf_outputs.items()
                    for index in indices
                ]
                new_body = relax.SeqExpr(
                    [*new_body.blocks, relax.BindingBlock(new_bindings)], new_body.body
                )

            new_body = LazyOutputMutator(self, self.mod).visit_expr(new_body)

        # Step 4. Add parameters of get_item and set_item (except index) to the function.
        params = [
            *func.params[:num_input],
            *self.extra_get_item_params,
            *self.extra_set_item_params,
        ]

        # Step 5. Find all shape parameters that should be retained as
        # parameters.
        symbolic_vars = relax.analysis.defined_symbolic_vars(func)
        if symbolic_vars:

            def unpack_ty(ty):
                if isinstance(ty, relax.TupleType):
                    for field in ty.fields:
                        yield from unpack_ty(field)
                else:
                    yield ty

            # direct iterate over the type annotation
            for param in func.params[num_input:]:
                for ty in unpack_ty(param.ty):
                    if isinstance(ty, tvm.ir.PrimType | relax.ShapeType):
                        params.append(relax.Var("symbolic_var_holder", ty))

        return relax.Function(
            params,
            new_body,
            relax.AnyType(),
            attrs=func.attrs,
            is_pure=False,
        ).without_attr("relax.force_pure")


@mutator
class LazyInputMutator(PyExprMutator):
    def __init__(self, func_creator, mod: IRModule | None = None) -> None:
        self.func_creator = func_creator
        super().__init__(mod)

    def visit_function_(self, func: relax.Function) -> relax.Expr:
        if "num_input" in func.attrs:
            num_input = func.attrs["num_input"]
        else:
            num_input = 0

        params = list(func.params)[num_input:]
        if len(params) == 1 and isinstance(params[0].ty, relax.TupleType):
            self.tuple_param = params[0]
            self.params = {}
        else:
            self.tuple_param = None
            self.params = {var: i for i, var in enumerate(params)}
        func = relax.Function(
            func.params[:num_input],
            func.body,
            func.ret_ty,
            is_pure=False,
            attrs=func.attrs,
            span=func.span,
        ).without_attr("relax.force_pure")
        output = super().visit_function_(func)
        self.tuple_param = None
        self.params = {}
        return output

    def visit_var_(self, var: relax.Var) -> relax.Expr:
        if var in self.params:
            index = self.params[var]
            get_item_result = self.builder_.emit(
                relax.Call(
                    relax.ExternFunc(self.func_creator.fget_item),
                    self.func_creator.extra_get_item_params + [relax.prim_value(index)],
                    None,
                    [relax.AnyType()],
                )
            )
            match_cast = relax.MatchCast(var, get_item_result, var.ty)
            self.builder_.emit_normalized(match_cast)

            del self.params[var]

        return super().visit_var_(var)

    def visit_tuple_getitem_(self, node: relax.TupleGetItem) -> relax.Expr:
        ty = node.ty

        node = super().visit_tuple_getitem_(node)

        if self.tuple_param is not None and node.tuple_value.same_as(self.tuple_param):
            get_item_result = self.builder_.emit(
                relax.Call(
                    relax.ExternFunc(self.func_creator.fget_item),
                    self.func_creator.extra_get_item_params + [relax.prim_value(node.index)],
                    None,
                    [relax.AnyType()],
                )
            )
            return self.builder_.match_cast(get_item_result, ty)
        else:
            return node


@mutator
class LazyOutputMutator(PyExprMutator):
    def __init__(self, func_creator, mod: IRModule | None = None) -> None:
        self.func_creator = func_creator
        self.killed_vars = set()
        super().__init__(mod)

    def visit_var_(self, var: relax.Var) -> None:
        assert var not in self.killed_vars
        return super().visit_var_(var)

    def visit_var_binding_(self, binding: relax.VarBinding) -> None:
        if binding.var == self.func_creator.out_tuple_var:
            # The function after rewriting returns a empty tuple.
            func_output = self.builder_.emit(relax.Tuple([]))
            self.set_var_remap(binding.var, func_output)
            return

        super().visit_var_binding_(binding)

        if binding.var in self.func_creator.memory_free_insertion:
            for var in self.func_creator.memory_free_insertion[binding.var]:
                if var in self.func_creator.out_tuple_map:
                    self.killed_vars.add(var)
                    for index in self.func_creator.out_tuple_map[var]:
                        # rewrite set item
                        self.builder_.emit(
                            relax.Call(
                                relax.ExternFunc(self.func_creator.fset_item),
                                self.func_creator.extra_set_item_params
                                + [index, super().visit_var_(var)],
                                None,
                                [relax.AnyType()],
                            ),
                            name_hint="_",
                        )

                if var in self.func_creator.input_params_set:
                    self.builder_.emit(
                        relax.op.vm.kill_object(super().visit_var_(var)), name_hint="_"
                    )


@tvm.transform.module_pass(opt_level=0, name="LazyTransformParams")
class LazyTransformParams:
    """
    Convert transform_params functions into a lazy version.
    (Load the input to memory on demand, and immediately free it after the last use.)

    Note: ToNonDataflow() and RemovePurityTracking() should be invoked before this pass.

    Parameters
    ----------
    fget_item: str
        The name of the get_item function.
    fset_item: str
        The name of the set_item function.
    extra_get_item_params: list of relax.Var
        The parameters of the get_item function except index.
        The given parameters will be placed before index.
        For example, if extra_get_item_params is [param1, param2], then the pass will generate
        call_packed(fget_item, [param1, param2, index])
    extra_set_item_params: list of relax.Var
        The parameters of the set_item function except index and value.
        The given parameters will be placed before index and value.
        For example, if extra_set_item_params is [param1, param2], then the pass will generate
        call_packed(fset_item, [param1, param2, index, value])
    """

    def __init__(
        self,
        fget_item="get_item",
        fset_item="set_item",
        extra_get_item_params=None,
        extra_set_item_params=None,
    ) -> None:
        self.fget_item = fget_item
        self.extra_get_item_params = [] if extra_get_item_params is None else extra_get_item_params
        assert self.fget_item is not None, "transforming set_item only is not supported"
        self.fset_item = fset_item
        self.extra_set_item_params = [] if extra_set_item_params is None else extra_set_item_params

    def transform_module(self, mod: IRModule, ctx: tvm.transform.PassContext) -> IRModule:
        lazy_mutator = LazyTransformParamsFuncCreator(
            self.fget_item,
            self.fset_item,
            self.extra_get_item_params,
            self.extra_set_item_params,
            mod,
        )
        builder = relax.BlockBuilder(mod)
        for gv, _ in mod.functions_items():
            if gv.name_hint.endswith("transform_params"):
                func = mod[gv]
                if not isinstance(func, relax.Function):
                    continue
                func = lazy_mutator.transform(func)
                builder.update_func(gv, func)

        return builder.get()
