# 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
# pylint: disable=redefined-builtin
# ruff: noqa: F821
"""The base Relax operators."""

from collections.abc import Callable

import tvm_ffi

import tvm
import tvm.runtime
from tvm.ir import Call
from tvm.runtime import Object, ObjectConvertible

from ..expr import Expr, ExternFunc, GlobalVar, ShapeExpr, StringImm, Var
from ..type import TensorType, Type
from ..utils import convert_to_expr
from . import _ffi_api

py_print = print  # pylint: disable=invalid-name


def register_gradient(
    op_name: str,
    fgradient: Callable[[Var, Call, Var, "BlockBuilder"], list[Expr]] | None = None,
    level: int = 10,
):
    """Register operator gradient function for a relax operator.

    Parameters
    ----------
    op_name: str
        The name of the op.

    fgradient: function (orig_var: Var, orig_call: Call, output_grad: Var, ctx: BlockBuilder)
         -> partials: List[Expr]
        The gradient function being used.

    level: int
        The priority level
    """
    return tvm.ir.register_op_attr(op_name, "FPrimalGradient", fgradient, level)


def null_value() -> Call:
    """Create a call node that represents a null value object.

    Returns
    -------
    ret: Call
        The created call node.
    """
    return _ffi_api.null_value()  # type: ignore


def _wrap_inline_arg_tuple(args) -> Expr:
    """Helper function to wrap argument tuple

    Normalize the arguments provided the functions that accept a tuple
    of arguments, and require the tuple of arguments to be written
    in-line.  If the arguments provided are a single relax expression,
    and are not a reference to a relax tuple, then wrap them into an
    in-line relax Tuple.

    """
    if isinstance(args, tuple | list):
        return tvm.relax.Tuple([convert_to_expr(a) for a in args])
    elif (
        isinstance(args, Expr)
        and not isinstance(args, tvm.relax.Tuple)
        and (args.ty is None or not isinstance(args.ty, tvm.relax.TupleType))
    ):
        return tvm.relax.Tuple([args])
    else:
        return args


def call_tir(
    gvar: GlobalVar,
    args: Expr,
    out_ty: TensorType | list[TensorType],
    tir_vars: ShapeExpr | tuple[Expr] | list[Expr] | None = None,
) -> Call:
    """
    Call a tirx.prim_func and return the output.

    Parameters
    ----------
    gvar : GlobalVar
        The GlobalVar referring to a tirx PrimFunc.

    args : Expr
        The input arguments.

    out_ty : Union[TensorType, List[TensorType]]
        The type information of the call_tir output.
        It should be a single or a list of TensorType. Each one denotes the
        type information of a returned tensor.

    tir_vars : Optional[Union[ShapeExpr, Tuple[Expr], List[Expr]]]
        ShapeExpr representing a tuple of integers to unpack when calling func. Is null if not used

    Returns
    -------
    ret: Call
        A call node for the call_tir operator.
    """
    args = _wrap_inline_arg_tuple(args)

    if not isinstance(out_ty, list):
        out_ty = [out_ty]

    if isinstance(tir_vars, list | tuple):
        tir_vars = ShapeExpr(tir_vars)

    return _ffi_api.call_tir(gvar, args, out_ty, tir_vars)  # type: ignore


def call_tir_with_grad(
    gvar: GlobalVar,
    args: Expr,
    out_ty: TensorType | list[TensorType],
    te_grad_name: str,
    te_grad_kwargs: dict[str, Object] | None = None,
    tir_vars: ShapeExpr | tuple[Expr] | list[Expr] | None = None,
) -> Call:
    """
    Call a tirx.prim_func and return the output. This intrinsic will bind a te gradient function
    (refered by te_grad_name) to the call_tir_with_grad node. The te gradient function will be
    called by the Gradient pass.

    Parameters
    ----------
    gvar : GlobalVar
        The GlobalVar referring to a tirx PrimFunc.

    args : Expr
        The input arguments.

    out_ty : Union[TensorType, List[TensorType]]
        The type information of the call_tir_with_grad output.
        It should be a single or a list of TensorType. Each one denotes the
        type information of a returned tensor.

    te_grad_name : str
        The registered name of the te gradient function associated with the call_tir_with_grad
        node. Must be provided as a keyword argument.

    te_grad_kwargs : Dict[str, Object], optional
        The keyword arguments passed to the te gradient function.
        Optionally provided as a keyword argument. Default: {}.

    tir_vars : Optional[Union[ShapeExpr, Tuple[Expr], List[Expr]]]
        ShapeExpr representing a tuple of integers to unpack when calling func. Is null if not used

    Returns
    -------
    ret: Call
        A call node for the call_tir_with_grad operator.
    """
    args = _wrap_inline_arg_tuple(args)

    if not isinstance(out_ty, list):
        out_ty = [out_ty]

    if isinstance(tir_vars, list | tuple):
        tir_vars = ShapeExpr(tir_vars)

    if te_grad_kwargs is None:
        te_grad_kwargs = {}

    return _ffi_api.call_tir_with_grad(  # type: ignore
        gvar, args, out_ty, te_grad_name, te_grad_kwargs, tir_vars
    )


def call_tir_inplace(
    gvar: GlobalVar,
    args: Expr,
    inplace_indices: int | list[int],
    out_ty: TensorType | list[TensorType],
    tir_vars: ShapeExpr | tuple[Expr] | list[Expr] | None = None,
) -> Call:
    """
    Call a TIR PrimFunc and return the result, doing the specified computations in-place
    (based on the `inplace_indices` argument; outputs will alias the inputs
    selected by in-place indices).

    Warning: This operator is considered pure by the type system but actually mutates
    the arguments specified by `inplace_indices`. This operator should not be used directly,
    but rather should be inserted by passes that have checked whether it is safe to perform
    operations in-place (i.e., none of the arguments specified as an output is aliased or is
    live after calling call_tir_inplace).

    Direct calls to this operator should be done for testing purposes only.

    Parameters
    ----------
    gvar : GlobalVar
        The GlobalVar referring to a TIR PrimFunc.

    args : Expr
        The input arguments.

    inplace_indices : Union[int, List[int]]
        Specify which arguments should be used for in-place computations.
        If `inplace_indices` is a single integer, it will be made into a singleton list.
        Suppose `inplace_indices[i] = j`, where `j >= 0`. Then the `i`th output
        will be an alias of `args[j]`.
        If `inplace_indices[i] = -1`, then the `i`th output will be a freshly allocated tensor.
        At least one member of `inplace_indices` must not be -1.

    out_ty : Union[TensorType, List[TensorType]]
        The type information of the call_tir_inplace output.
        It should be a single `TensorType` or a list of `TensorType`.
        Each one denotes the type information of a returned tensor.
        If a list of `TensorType` is given, the result will be a tuple of `TensorType`.

    tir_vars : Optional[Union[ShapeExpr, Tuple[Expr], List[Expr]]]
        ShapeExpr representing a tuple of integers to unpack when calling func. Is null if not used

    Returns
    -------
    ret: Call
        A call node for the call_tir operator.
    """
    args = _wrap_inline_arg_tuple(args)

    if not isinstance(inplace_indices, list):
        inplace_indices = [inplace_indices]

    if not isinstance(out_ty, list):
        out_ty = [out_ty]

    if isinstance(tir_vars, list | tuple):
        tir_vars = ShapeExpr(tir_vars)

    return _ffi_api.call_tir_inplace(  # type: ignore
        gvar,
        args,
        inplace_indices,
        out_ty,
        tir_vars,
    )


def call_dps_packed(
    func: str | Expr,
    args: Expr,
    out_ty: TensorType | list[TensorType],
) -> Call:
    """
    Call a destination-passing-style packed function and return the output.

    Note: The called function is assumed to be _pure_ (other than modifying the designated
    output arguments). If the function _does_ result in other side effects, then the compiler
    may end up removing, reordering, or repeating those effects--no guarantees can be made.

    Parameters
    ----------
    func : Union[str, Expr]
        The destination-passing-style function, can be ExternFunc.

    args : Expr
        The input arguments.

    out_ty : Union[TensorType, List[TensorType]]
        The type information of the call_dps_packed output.
        It should be a single or a list of TensorType. Each one denotes the
        type information of a returned tensor.

    Returns
    -------
    ret: Call
        A call node for the call_dps_packed operator.
    """
    if isinstance(func, str):
        func = ExternFunc(func)

    args = _wrap_inline_arg_tuple(args)

    if not isinstance(out_ty, list):
        out_ty = [out_ty]

    return _ffi_api.call_dps_packed(func, args, out_ty)  # type: ignore


def call_py_func(
    func_name: str,
    args: Expr,
    out_ty: TensorType | list[TensorType],
) -> Call:
    """
    Call a Python function and return the output.

    Parameters
    ----------
    func_name : str
        The name of the Python function to call. This should correspond to a function
        in the IRModule's pyfuncs attribute.

    args : Expr
        The input arguments.

    out_ty : Union[TensorType, List[TensorType]]
        The type information of the call_py_func output.
        It should be a single or a list of TensorType. Each one denotes the
        type information of a returned tensor.

    Returns
    -------
    ret: Call
        A call node for the call_py_func operator.
    """
    args = _wrap_inline_arg_tuple(args)

    if not isinstance(out_ty, list):
        out_ty = [out_ty]

    return _ffi_api.call_py_func(func_name, args, out_ty)  # type: ignore


def call_builtin_with_ctx(
    func: str | Expr,
    args: Expr,
    *,
    ty_args: Type | list[Type] | None = None,
) -> Call:
    """Call a builtin function func.

    Parameters
    ----------
    func : Expr
        The builtin function to be called.

    args : Expr
        The input arguments.

    ty_args: Optional[Union[Type, List[Type]]]
        The type arguments to the call node.

    Returns
    -------
    ret: Call
        The created call node.
    """
    if isinstance(func, str):
        func = ExternFunc(func)

    args = _wrap_inline_arg_tuple(args)

    if ty_args is not None and not isinstance(ty_args, list | tuple):
        ty_args = [ty_args]

    return _ffi_api.call_builtin_with_ctx(  # type: ignore
        func,
        args,
        ty_args,  # type: ignore
    )


def make_closure(
    func: Expr,
    args: Expr,
) -> Object:
    """
    Create a closure with free variables and return the closure.

    Parameters
    ----------
    func : Expr
        The closure, can be ExternFunc or PrimFunc.

    args : Expr
        The input arguments.


    Returns
    -------
    ret: Object
        The VMClosure.
    """

    args = _wrap_inline_arg_tuple(args)

    return _ffi_api.make_closure(func, args)  # type: ignore


def invoke_closure(
    closure: Expr,
    args: Expr,
    ty_args: list[Type] | Type,
) -> Call:
    """
    Invoke a closure.

    Parameters
    ----------
    closure : Expr
        The VMClosure object.

    args : Expr
        The input arguments.

    type_args: Union[List[Type], Type]
        The type information arguments of the CallNode

    Returns
    -------
    ret: Call
        A call to `invoke_closure`.
    """
    args = _wrap_inline_arg_tuple(args)

    if not isinstance(ty_args, list | tuple):
        ty_args = [ty_args]

    return _ffi_api.invoke_closure(closure, args, ty_args)  # type: ignore


def render_object(val: tvm.Object) -> str:
    """
    Given a TVM Object, renders it in string form. Used for Relax printing and assertions.

    Parameters
    ----------
    val: tvm.Object
        An object to render

    Returns
    -------
    ret: str
        A string representing the value, ideally human-readable
    """
    if isinstance(val, tvm.runtime.Tensor):
        return str(val)
    if isinstance(val, tvm_ffi.Array):
        fields = ", ".join([render_object(val[i]) for i in range(len(val))])
        return f"({fields})"
    return str(val)


@tvm.register_global_func("relax.run.shape_to_tensor")
def relax_shape_to_tensor(shape_tuple: tvm_ffi.Shape) -> tvm.runtime.Tensor:
    """
    Takes a Shape and convert it to Tensor.

    Parameters
    ----------
    shape_tuple: tvm_ffi.Shape
        Shape tuple that we want to convert to Tensor at runtime
    """
    return tvm.runtime.tensor([int(v) for v in shape_tuple])


@tvm.register_global_func("relax.run.print")
def relax_print(format_str: str, *format_args: tvm.Object) -> None:
    """
    Takes a list of values to print, formats with the given format string.
    If the format string is empty, simply prints.

    Call from TVM script like this:
    `relax.print(value1, value2, ..., valueN, format=format_str)`
    or
    `relax.print(value1, value2, ..., valueN) # format_str defaults to ""`

    Parameters
    ----------
    format_str: str
        The last argument is a Python-style format string for printing the value

    format_args: List[Object]
        The values to print.
    """
    val_strs = map(render_object, format_args)
    if format_str == "":
        py_print(*val_strs)
    else:
        py_print(format_str.format(*val_strs))


def print(*values: list[Expr], format: str | Expr = "") -> Expr:
    """Print op to print the values

    Parameters
    ----------
    values : List[Expr]
        The values to print.

    format: Union[str, Expr]
        The format string or StringImm.

    Returns
    -------
    result : Expr
        A relax Call, which will print the value during runtime.
    """
    if isinstance(format, str):
        format = StringImm(format)

    return _ffi_api.print(values, format)  # type: ignore # pylint: disable=no-member


@tvm.register_global_func("relax.run.assert_op")
def relax_assert_op(condition: tvm.Object, format_str: str, *format_args: tvm.Object) -> None:
    """
    A variadic function. The first value serves as the assertion condition:
    If the condition is true, then the operator does nothing.
    If the condition is false, then the operator raises an assertion error.

    Arguments after the first value serve as format arguments for the error message;
    the last argument must be a format string for the error message (empty by default).
    If the format string is the empty string, then the error message will simply include
    a comma-separated list of the format arguments.
    The condition argument is not included in the format string.

    Parameters
    ----------
    condition: tvm.Object
        The assertion condition. Must be a boolean scalar.

    format_str: str
        The last argument is a Python-style format string for printing the value

    format_args: List[tvm.Object]
        Values used for formatting the string.
    """
    if not isinstance(format_str, str):
        raise ValueError(
            f"The format string argument to assert must be a string, given {type(format_str)})"
        )

    if isinstance(condition, bool | int):
        val = condition
    elif isinstance(condition, tvm.runtime.Tensor):
        # may happen if the original program had unknown shape or dtype for the tensor's type
        dtype = condition.dtype
        if dtype != "bool":
            raise ValueError(f"The condition must be a bool scalar, but given a {dtype} tensor")
        shape = condition.shape
        if len(shape) != 0:
            raise ValueError(f"The condition must be a scalar, but it has a shape of {shape}")

        val = condition.numpy()

    else:
        # should be guaranteed by the type system
        raise ValueError(
            f"The condition for relax assert must be a bool, int, or Tensor, "
            f"but received a {type(condition)}."
        )

    if not val:
        error_message = "Assertion Failed"
        if format_args or format_str != "":
            rendered = map(render_object, format_args)
            if format_str != "":
                error_message = format_str.format(*rendered)
            else:
                error_message = ", ".join(rendered)
        raise AssertionError(error_message)


def assert_op(
    condition: Expr,
    format_args: Expr | list[Expr] | None = None,
    format: str | Expr = "",
) -> Expr:
    """
    Create a call to Relax's assert_op operation (`assert` is reserved in Python,
    so the name must be distinct).

    Parameters
    ----------
    condition: Expr
        The assertion condition.

    format_args: Optional[Union[Expr, List[Expr]]]
        Format arguments for the error message if the condition fails.

    format: Union[str, Expr]
        The format string or StringImm for the error message.

    Returns
    -------
    result : Expr
        A Call to the Relax assert operation.
    """
    if not isinstance(condition, Expr):
        condition = tvm.relax.prim_value(condition)

    if format_args is None:
        format_args = []
    elif isinstance(format_args, Expr):
        format_args = [format_args]

    if isinstance(format, str):
        format = StringImm(format)

    return _ffi_api.assert_op(condition, format_args, format)  # type: ignore


def shape_of(expr: Expr) -> Expr:
    """Get shape of a tensor.

    Parameters
    ----------
    expr : Expr
        The input Expr.

    Returns
    -------
    result : Expr
        A relax Call, which gets the shape of the input
    """
    return _ffi_api.shape_of(expr)  # type: ignore # pylint: disable=no-member


def size(expr: Expr) -> Expr:
    """Get the total number of elements in a tensor.

    Parameters
    ----------
    expr : Expr
        The input tensor.

    Returns
    -------
    result : Expr
        A scalar tensor of dtype int64 containing the total number of elements.
    """
    return _ffi_api.size(expr)  # type: ignore # pylint: disable=no-member


def tensor_to_shape(expr: Expr) -> Expr:
    """Convert tensor to shape expr.
    Parameters
    ----------
    expr : Expr
        The input Expr
    Returns
    -------
    result : Expr
        A relax Call, which transforms the tensor values to the shape
    """
    return _ffi_api.tensor_to_shape(expr)  # type: ignore # pylint: disable=no-member


def shape_to_tensor(expr: Expr) -> Expr:
    """Convert shape to tensor expr.
    Parameters
    ----------
    expr : Expr
        The input Expr
    Returns
    -------
    result : Expr
        A relax Call, which transforms the shape values to the tensor
    """
    return _ffi_api.shape_to_tensor(expr)  # type: ignore # pylint: disable=no-member


def call_inplace_packed(
    func: str | ExternFunc | GlobalVar,
    *args: Expr,
    inplace_indices: int | list[int],
    ty_args: Type | list[Type],
) -> Expr:
    """
    Construct a call to a packed function that consumes some of its arguments "in-place"
    and returns the mutated arguments (aliased), but should be considered to be otherwise pure.
    The `inplace_indices` argument indicates which of the outputs are mutated arguments.

    The resulting call will have the same semantics as calling the packed function directly.

    Note: This should be used for cases when the user knows that calling the packed function
    with these arguments will **in reality** not cause any other side effects.
    If it is used for a call that **does** result in other side effects, then the compiler
    may end up removing, reordering, or repeating that call, with no guarantees
    made about any side effects from the callee.

    Warning: This operator as treated as pure by the type system even though it *is* performing
    side effects (mutating some arguments). It is therefore incumbent upon the user to ensure
    that it is being used safely (viz., that mutated arguments are not live after the mutation,
    that they do not alias values live after the mutation).

    Parameters
    ----------
    func : Union[str, ExternFunc]
      The name (global symbol) for a PackedFunc or an ExternFunc node.

    args: Expr
      The arguments for the PackedFunc.

    inplace_indices : Union[int, List[int]]
      Specify which arguments should be used for in-place computations.
      If `inplace_indices` is a single integer, it will be made into a singleton list.
      Suppose `inplace_indices[i] = j`, where `j >= 0`. Then the `i`th output
      will be an alias of `args[j]`.
      If `inplace_indices[i] = -1`, then the `i`th output will be a freshly allocated tensor.
      At least one member of `inplace_indices` must not be -1.

    ty_args: Union[Type, List[Type]]
        The list of type information arguments (giving the type information for the returned value).

    Returns
    -------
    result : Expr
      A Relax call, corresponding to
      `call_pure_packed(ExternFunc(func), args, DictAttrs(kwargs), ty_args)`
    """
    if isinstance(func, ExternFunc):
        func = func.global_symbol

    op = ExternFunc(func)
    args = tuple(convert_to_expr(a) for a in args)
    if ty_args is None:
        raise ValueError("R.call_pure_packed is required to have type_args")
    if isinstance(ty_args, tuple):  # type: ignore
        ty_args = list(ty_args)
    elif not isinstance(ty_args, list):
        ty_args = [ty_args]
    if not isinstance(inplace_indices, list):
        inplace_indices = [inplace_indices]

    return _ffi_api.call_inplace_packed(op, args, inplace_indices, ty_args)  # type: ignore # pylint: disable=no-member


def call_pure_packed(
    func: str | ExternFunc | GlobalVar,
    *args: Expr,
    ty_args: Type | list[Type],
) -> Expr:
    """
    Construct a call to a packed function that should be treated as pure,
    even though packed calls are normally not treated as pure.

    The resulting call will have the same semantics as calling the packed function directly.

    Note: This should be used for cases when the user knows that calling the packed function
    with these arguments will **in reality** not cause any side effects.
    If it is used for a call that **does** result in side effects, then the compiler
    may end up removing, reordering, or repeating that call, with no guarantees
    made about any side effects from the callee.

    Parameters
    ----------
    func : Union[str, ExternFunc]
      The name (global symbol) for a PackedFunc or an ExternFunc node.

    args: Expr
      The arguments for the PackedFunc.

    ty_args: Union[Type, List[Type]]
        The list of type information arguments (giving the type information for the returned value).

    Returns
    -------
    result : Expr
      A Relax call, corresponding to
      `call_pure_packed(ExternFunc(func), args, DictAttrs(kwargs), ty_args)`
    """
    if isinstance(func, ExternFunc):
        func = func.global_symbol

    op = ExternFunc(func)
    args = tuple(convert_to_expr(a) for a in args)

    if ty_args is None:
        raise ValueError("R.call_pure_packed is required to have type_args")

    if isinstance(ty_args, tuple):  # type: ignore
        ty_args = list(ty_args)
    elif not isinstance(ty_args, list):
        ty_args = [ty_args]

    ty_args = [
        (ty() if callable(ty) else ty.asobject() if isinstance(ty, ObjectConvertible) else ty)
        for ty in ty_args
    ]

    # note: if we need attributes, we can also take them here

    return _ffi_api.call_pure_packed(op, args, None, ty_args)  # type: ignore # pylint: disable=no-member


def invoke_pure_closure(
    closure: Expr,
    args: Expr,
    ty_args: list[Type] | Type,
) -> Call:
    """
    Invoke a closure and indicate to the compiler that it is pure.

    Note: This should be used for cases when the user knows that calling the closure
    with these arguments will **in reality** not cause any side effects.
    If it is used for a call that _does_ result in side effects, then the compiler
    may end up removing, reordering, or repeating that call, with no guarantees
    made about any side effects from the callee.

    Parameters
    ----------
    closure : Expr
        The VMClosure object.

    args : Expr
        The input arguments.

    type_args: Union[List[Type], Type]
        The type information arguments of the CallNode

    Returns
    -------
    ret: Call
        A call to `invoke_pure_closure`.
    """
    args = _wrap_inline_arg_tuple(args)

    if not isinstance(ty_args, list | tuple):
        ty_args = [ty_args]

    return _ffi_api.invoke_pure_closure(closure, args, ty_args)  # type: ignore


def to_vdevice(data, dst_vdevice) -> Expr:
    """Copy data to the destination device. This
    operator helps data transferring between difference devices for
    heterogeneous execution.

    Parameters
    ----------
    data : Expr
        The tensor to be copied.

    dst_device : VDevice
        The destination device where the data is copied to.

    Returns
    -------
    result : Expr
        The copied result.
    """
    return _ffi_api.to_vdevice(data, dst_vdevice)  # type: ignore


def hint_on_device(data, dst_vdevice, memory_scope="global") -> Expr:
    """It provides a hint specifying the device on which the input data should be executed.
    This hint is utilized by RealizeVDevice to propagate the virtual device."

    Parameters
    ----------
    data : Expr
        The tensor to be copied.

    dst_device : Device
        The destination device where the data is supposed to be executed.

    memory_scope: String
       Memory scope of buffer on target device.

    Returns
    -------
    result : Expr
        The result.
    """
    return _ffi_api.hint_on_device(data, dst_vdevice, memory_scope)  # type: ignore
