# 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=abstract-method,unused-argument
# pylint: disable=missing-function-docstring,missing-module-docstring

import tvm
from tvm.ir import Call, Op
from tvm.ir.module import IRModule

from ..expr import Expr, Function, ShapeExpr
from ..expr_functor import PyExprVisitor, visitor


def estimate_memory_usage(mod: IRModule | Function) -> str:
    """Analysis function that estimates the memory usage of Relax functions
    in an IRModule. The estimation includes the total memory size needed to
    be allocated before and after memory planning.

    The result might be over-estimated, as the estimation is static, which
    does not consider control flows (such as "if" and cross-function calls).
    It simply accumulates the size of every alloc_tensor and alloc_storage.

    This analysis function is used to demonstrate the effect of memory
    planning.

    Parameters
    ----------
    mod : Union[IRModule, Function]
        The input IRModule whose functions inside are to be analyzed.
        If the input is a Function, we will wrap it with a IRModule, with
        the function named "main".

    Returns
    -------
    est : str
        The estimation information, in the form of a string.

    Notes
    -----
    We regards "relax.memory.alloc_tensor/storage" as the results produced by memory planning.
    """

    @visitor
    class MemoryEstimator(PyExprVisitor):
        """The IR visitor which estimates the memory usage of each Relax function.

        Attributes
        ----------
        total_alloc_tensor_mem : int
            The total memory size of alloc_tensor, in bytes.

        total_const_size_tensor_num : int
            The number of constant-size tensors.

        total_dyn_size_tensor_num : int
            The number of dynamic-size tensors.

        planned_alloc_mem : int
            The total memory size of memory.alloc_storage after memory planning, in bytes.

        planned_mem_num : int
            The number of memory.alloc_storages.
        """

        total_alloc_tensor_mem: int
        total_const_size_tensor_num: int
        total_dyn_size_tensor_num: int
        planned_alloc_mem: int
        planned_mem_num: int
        builtin_alloc_tensor_op = Op.get("relax.builtin.alloc_tensor")
        memory_alloc_tensor_op = Op.get("relax.memory.alloc_tensor")
        memory_alloc_storage_op = Op.get("relax.memory.alloc_storage")

        def estimate(self, mod: IRModule) -> str:
            estimation: str = ""
            for global_var, func in mod.functions_items():
                if not isinstance(func, Function):
                    continue

                self.cleanup()
                self.visit_expr(func)
                estimation += self.generate_est_string(global_var.name_hint) + "\n"

            if estimation != "":
                estimation = "Memory usage estimation:\n" + estimation
            return estimation

        def cleanup(self) -> None:
            self.total_alloc_tensor_mem = 0
            self.total_const_size_tensor_num = 0
            self.total_dyn_size_tensor_num = 0
            self.planned_alloc_mem = 0
            self.planned_mem_num = 0

        def visit_call_(self, call: Call) -> None:  # pylint: disable=arguments-differ
            if call.op == self.builtin_alloc_tensor_op:
                self.accumulate_builtin_tensor_alloc(
                    shape=call.args[0], dtype_str=call.args[1].value
                )
            elif call.op == self.memory_alloc_tensor_op:
                self.accumulate_tensor_alloc(shape=call.args[2], dtype_str=call.args[3].value)
            elif call.op == self.memory_alloc_storage_op:
                self.accumulate_storage_alloc(size=call.args[0])

        def calculate_size(self, shape: Expr, dtype_str: str) -> int:
            if not isinstance(shape, ShapeExpr):
                raise TypeError(
                    "The shape of relax.builtin.alloc_tensor and "
                    "relax.memory.alloc_tensor is expected to be ShapeExpr"
                )
            size: int = 1
            for dim_len in shape.values:
                if not isinstance(dim_len, tvm.tirx.IntImm):
                    self.total_dyn_size_tensor_num += 1
                    return -1
                size *= dim_len.value
            dtype = tvm.DataType(dtype_str)
            return size * ((dtype.bits + 7) // 8) * dtype.lanes

        def accumulate_builtin_tensor_alloc(self, shape: Expr, dtype_str: str) -> None:
            size = self.calculate_size(shape, dtype_str)
            if size == -1:
                return
            self.total_const_size_tensor_num += 1
            self.total_alloc_tensor_mem += size
            self.planned_mem_num += 1
            self.planned_alloc_mem += size

        def accumulate_tensor_alloc(self, shape: Expr, dtype_str: str) -> None:
            size = self.calculate_size(shape, dtype_str)
            if size == -1:
                return
            self.total_const_size_tensor_num += 1
            self.total_alloc_tensor_mem += size

        def accumulate_storage_alloc(self, size: Expr) -> None:
            if not isinstance(size, ShapeExpr):
                raise TypeError(
                    "The size of relax.memory.alloc_storage is expected to be ShapeExpr"
                )

            self.planned_mem_num += 1
            self.planned_alloc_mem += size.values[0].value

        def generate_est_string(self, func_name: str) -> str:
            est = (
                f" * Without memory planning, there are {self.total_const_size_tensor_num} "
                "constant-size memory allocation(s) with total size "
                "{0:.4} GB".format(self.total_alloc_tensor_mem / 2**30)
            )
            if self.total_dyn_size_tensor_num > 0:
                est += f", and {self.total_dyn_size_tensor_num} dynamic-size allocation(s)"
            est += (
                f".\n * With memory planning, there are {self.planned_mem_num} constant-size "
                "memory allocation(s) with total size "
                "{0:.4} GB.\n".format(self.planned_alloc_mem / 2**30)
            )
            if self.total_alloc_tensor_mem != 0:
                est += (
                    " * Memory planning reduces constant memory size to "
                    f"{self.planned_alloc_mem / self.total_alloc_tensor_mem:.1%}."
                )
            return "- Function " + func_name + ":\n" + est

    if isinstance(mod, Function):
        mod = tvm.IRModule({tvm.ir.GlobalVar("foo"): mod})

    return MemoryEstimator().estimate(mod)
