# 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=unrecognized-inline-option
"""Function data types."""

import collections
import inspect
from collections.abc import Callable, Mapping
from typing import Optional

import tvm_ffi

import tvm
import tvm.runtime
from tvm.ir import BaseFunc, Range
from tvm.runtime import Object, Scriptable

from ..runtime._tensor import Tensor
from . import _ffi_api
from .buffer import Buffer
from .expr import Expr, Var


@tvm_ffi.register_object("tirx.PrimFunc")
class PrimFunc(BaseFunc, Scriptable):
    """A function declaration expression.

    Parameters
    ----------
    params: List[Union[tvm.tirx.Var, tvm.tirx.Buffer]]
        List of input parameters to the function.

    body: tvm.tirx.Stmt
        The body of the function.

    ret_type: tvm.ir.Type
        The return type annotation of the function.

    buffer_map : Map[tvm.tirx.Var, tvm.tirx.Buffer]
        The buffer binding map.

    attrs: Optional[tvm.Attrs]
        Attributes of the function, can be None

    span : Optional[Span]
        The location of this itervar in the source code.
    """

    def __init__(self, params, body, ret_type=None, buffer_map=None, attrs=None, span=None):
        # Legacy compatibility: expand body-carrying leaf stmt wrappers
        # (e.g. DeclBuffer/AllocBuffer forms) into SeqStmt form.
        from .stmt import _normalize_legacy_stmt

        body = _normalize_legacy_stmt(body)
        if ret_type is None:
            ret_type = tvm.ir.Type.missing()
        param_list = []
        buffer_map = {} if buffer_map is None else buffer_map
        for x in params:
            x = tvm.runtime.convert(x) if not isinstance(x, Object) else x
            if isinstance(x, Buffer):
                var = Var(x.name, dtype="handle")
                param_list.append(var)
                buffer_map[var] = x
            elif isinstance(x, Var):
                param_list.append(x)
            else:
                raise TypeError("params can only contain Var or Buffer")

        if attrs is None:
            attrs = tvm.ir.make_node("ir.DictAttrs")

        self.__init_handle_by_constructor__(
            _ffi_api.PrimFunc,
            param_list,
            body,
            ret_type,
            buffer_map,
            attrs,
            span,
        )  # type: ignore

    def with_body(self, new_body, span=None):
        """Create a new PrimFunc with the same set signatures but a new body.

        Parameters
        ----------
        new_body : Stmt
            The new body.

        span : Optional[Span]
            The location of this itervar in the source code.

        Returns
        -------
        new_func : PrimFunc
            The created new function.
        """
        return PrimFunc(
            self.params,
            new_body,
            self.ret_type,
            self.buffer_map,
            self.attrs,
            span,
        )

    def specialize(self, param_map: Mapping[Var, Expr | Buffer]):
        """Specialize parameters of PrimFunc

        Parameters
        ----------

        param_map : Mapping[Var, Union[Expr, Buffer]]
            The mapping from function params to the instance

        Examples
        --------
        We can define a Meta TIR function with symbolic shape:

        .. code-block:: python

            @T.prim_func(s_tir=True)
            def mem_copy(a: T.handle, b: T.handle, m: T.int32, n: T.int32) -> None:
                A = T.match_buffer(a, (m, n), "float32")
                B = T.match_buffer(b, (m, n), "float32")

                for i, j in T.grid(m, n):
                    with T.sblock():
                        vi, vj = T.axis.remap("SS", [i, j])
                        B[vi, vj] = A[vi, vj]

        Then we can make it specialized with given shapes or buffers.

        .. code-block:: python

            a, _, m, n = mem_copy.params
            func = mem_copy.specialize({a: tirx.decl_buffer((16, 16))})
            # or
            func = mem_copy.specialize({n: 16, m: 16})

        The specialized function:

        .. code-block:: python

            @T.prim_func(s_tir=True)
            def mem_copy_16_16(a: T.handle, b: T.handle) -> None:
                A = T.match_buffer(a, (16, 16), "float32")
                B = T.match_buffer(b, (16, 16), "float32")

                for i, j in T.grid(16, 16):
                    with T.sblock():
                        vi, vj = T.axis.remap("SS", [i, j])
                        B[vi, vj] = A[vi, vj]

        Returns
        -------
        func : PrimFunc
            The new function with parameter specialized
        """
        return _ffi_api.Specialize(self, param_map)  # type: ignore


@tvm_ffi.register_object("tirx.TensorIntrin")
class TensorIntrin(Object):
    """A tensor intrinsic.

    Parameters
    ----------
    desc : PrimFunc
        The function to describe the computation.

    impl : PrimFunc
        The function of the implementation for the execution.
    """

    def __init__(self, desc, impl):
        self.__init_handle_by_constructor__(_ffi_api.TensorIntrin, desc, impl)

    @staticmethod
    def register(name: str, desc: PrimFunc, impl: PrimFunc, override: bool = False):
        """Register a tensor intrinsic with its name.

        Parameters
        ----------
        name : str
            The name of the TensorIntrin to register.
        desc : PrimFunc
            The function to describe the computation.
        impl : PrimFunc
            The function of the implementation for the execution.
        override: bool
            Whether override existing intrinsic.
        """
        return _ffi_api.TensorIntrinRegister(name, TensorIntrin(desc, impl), override)  # type: ignore

    @staticmethod
    def get(name: str, allow_missing: bool = False) -> Optional["TensorIntrin"]:
        """Look up a tensor intrinsic by its name.

        Parameters
        ----------
        name : str
            The name of the TensorIntrin to look up.

        allow_missing : bool
            Whether to allow missing tensor intrin. If False, raise an error if the tensor intrin
        doesn't exist.

        Returns
        -------
        result : Optional[TensorIntrin]
            The TensorIntrin with the specified name, or None if not found.
        """
        return _ffi_api.TensorIntrinGet(name, allow_missing)  # pylint: type: ignore


@tvm_ffi.register_object("tirx.IndexMap")
class IndexMap(Object):
    """A mapping from multi-dimensional indices to another set of multi-dimensional indices

    Parameters
    ----------
    initial_indices : List[Var]
        Variables representing the indices prior to remapping.
    final_indices : List[Expr]
        Expressions defining the indices after remapping.
    inverse_index_map : Union[Callable, Optional[IndexMap]]
        The optional pre-defined inverse index map.
        When this is defined, IndexMap::Inverse will return the pre-defined inverse index map.
        Otherwise, the inverse index map will be computed on the fly.
        It is the user's responsibility to ensure the correctness of the pre-defined inverse
        index map.
    """

    initial_indices: list[Var]
    final_indices: list[Expr]

    # Sentinel value used to indicate which groups of pre-flattening axes
    # should be used to post-flattening axes axes.  See
    # Stage.transform_layout for more details.
    AXIS_SEPARATOR = "axis_separator"

    def __init__(self, initial_indices, final_indices, inverse_index_map):
        if isinstance(inverse_index_map, Callable):
            inverse_index_map = IndexMap.from_func(inverse_index_map)
        self.__init_handle_by_constructor__(
            _ffi_api.IndexMap, initial_indices, final_indices, inverse_index_map
        )

    @staticmethod
    def from_func(
        mapping_function: Callable,
        ndim: int | None = None,
        inverse_index_map: Callable | Optional["IndexMap"] = None,
        *,
        index_dtype: str = "int64",
    ):
        """Create an index map from a function

        Parameters
        ----------
        mapping_function : Callable

            The function to map from source indices to target indices.
            The function should accept `tirx.Var` parameters and return
            a either a `tirx.Expr`, or a list of `tirx.Expr`.
            Returning a `tirx.Expr` is equivalent to returning a
            list of length 1 containing that `tirx.Expr`.

        ndim: Optional[int]

            The dimensionality of the buffer to which this
            transformation should be applied.  If mapping_function uses
            variadic argument `*args`, `ndim` must be specified.  If
            mapping_function does not use variadic arguments, ndim is
            optional.

        inverse_index_map : Union[Callable, Optional[IndexMap]]
            The optional pre-defined inverse index map.
            When this is defined, IndexMap::Inverse will return the pre-defined inverse index map.
            Otherwise, the inverse index map will be computed on the fly.
            It is the user's responsibility to ensure the correctness of the pre-defined inverse
            index map.

        Returns
        -------
        index_map: IndexMap

            Returns an IndexMap representing the `mapping_function`.

        """
        index_map, axis_separators = IndexMap.from_func_with_separators(
            mapping_function,
            ndim,
            inverse_index_map,
            index_dtype=index_dtype,
        )
        assert not axis_separators, (
            "The mapping_function provided to IndexMap.from_func "
            "may not return IndexMap.AXIS_SEPARATOR.  "
            "If required, please use IndexMap.from_func_with_separators instead."
        )
        return index_map

    @staticmethod
    def from_func_with_separators(
        mapping_function: Callable,
        ndim: int | None = None,
        inverse_index_map: Callable | Optional["IndexMap"] = None,
        *,
        index_dtype: str = "int64",
    ):
        """Create an index map from a function

        Parameters
        ----------
        mapping_function : Callable

            The function to map from source indices to target indices.
            The function should accept tirx.Var parameters and return
            either a `tirx.Expr` or a list.  Each element of the
            returned list should be either a `tirx.Expr` or the
            object `IndexMap.AXIS_SEPARATOR`.  Returning a
            `tirx.Expr` is equivalent to returning a list of length
            1 containing that `tirx.Expr`.

        ndim: Optional[int]

            The dimensionality of the buffer to which this
            transformation should be applied.  If mapping_function uses
            variadic argument `*args`, ndim must be specified.  If
            mapping_function does not use variadic arguments, ndim is
            optional.

        inverse_index_map : Union[Callable, Optional[IndexMap]]
            The optional pre-defined inverse index map.
            When this is defined, IndexMap::Inverse will return the pre-defined inverse index map.
            Otherwise, the inverse index map will be computed on the fly.
            It is the user's responsibility to ensure the correctness of the pre-defined inverse
            index map.

        index_dtype : str
            The default index dtype to use for input iters in the mapping function.

        Returns
        -------
        ret: Tuple[IndexMap, List[int]]

            Returns a tuple whose first element is an IndexMap
            representing the `mapping_function`, and whose second index
            is a list of indices at which `IndexMap.AXIS_SEPARATOR`
            occurred.

        """
        params = inspect.signature(mapping_function).parameters

        args = []
        var_arg_name = None
        kwargs = collections.OrderedDict()

        for name, param in params.items():
            if param.kind in [
                inspect.Parameter.POSITIONAL_ONLY,
                inspect.Parameter.POSITIONAL_OR_KEYWORD,
            ]:
                args.append(tvm.tirx.Var(name, index_dtype))

            elif param.kind == inspect.Parameter.VAR_POSITIONAL:
                var_arg_name = name

            elif param.kind == inspect.Parameter.KEYWORD_ONLY:
                kwargs[name] = tvm.tirx.Var(name, index_dtype)

            else:
                raise ValueError("transform_layout mapping may not have *args")

        # Now that all the named arguments have been collected,
        # everything that remains should go to the *args, if
        # specified.
        if var_arg_name is not None:
            assert ndim is not None, "ndim must be specified when *args is used"
            num_var_args = ndim - len(args) - len(kwargs)
            for i in range(num_var_args):
                args.append(tvm.tirx.Var(f"{var_arg_name}_{i}", index_dtype))

        mapping = mapping_function(*args, **kwargs)

        initial_indices = args + list(kwargs.values())

        final_indices = []
        axis_separators = []

        try:
            iter(mapping)
            is_iterable = True
        except TypeError:
            is_iterable = False

        if is_iterable:
            for val in mapping:
                if tvm.ir.is_prim_expr(val):
                    final_indices.append(val)
                elif val is IndexMap.AXIS_SEPARATOR:
                    axis_separators.append(len(final_indices))
                else:
                    raise TypeError(
                        "Expected mapping function to return list of "
                        "either tvm.ir.Expr or IndexMap.AXIS_SEPARATOR.  "
                        f"Instead received {val} of type {type(val)}."
                    )
        else:
            final_indices.append(mapping)

        return IndexMap(initial_indices, final_indices, inverse_index_map), axis_separators

    def is_equivalent_to(self, other_map: "IndexMap", analyzer=None) -> bool:
        """Return if the index maps are equivalent.

        Parameters
        ----------
        other_map: IndexMap

            The IndexMap to which the comparison should be made.

        analyzer : Optional[tvm.arith.Analyzer]

            The analyzer to use while comparing the mapped indices.  When
            provided, its accumulated bindings and constraints are reused so
            that maps that are only equivalent under those bindings can be
            proven equal.

        Returns
        -------
        is_equivalent: bool

            True if the two mappings represent the same
            transformation, otherwise False
        """
        if len(self.initial_indices) != len(other_map.initial_indices):
            return False
        if len(self.final_indices) != len(other_map.final_indices):
            return False

        if analyzer is None:
            analyzer = tvm.arith.Analyzer()

        mapped_other_final_indices = other_map.map_indices(self.initial_indices, analyzer=analyzer)
        for self_index, other_index in zip(self.final_indices, mapped_other_final_indices):
            if not analyzer.can_prove_equal(self_index, other_index):
                return False

        return True

    def map_indices(self, indices: list[Expr], analyzer=None) -> list[Expr]:
        """Apply the index map to a set of indices

        Parameters
        ----------
        indices : List[Expr]
            The indices to be mapped
        analyzer : Optional[tvm.arith.Analyzer]
            The analyzer to use while simplifying mapped indices.

        Returns
        -------
        result : List[Expr]
            The mapped indices
        """
        return _ffi_api.IndexMapMapIndices(self, indices, analyzer)

    def map_shape(self, shape: list[Expr], analyzer=None) -> list[Expr]:
        """Apply the index map to a buffer shape

        Parameters
        ----------
        shape : List[Expr]
            The buffer shape to be mapped
        analyzer : Optional[tvm.arith.Analyzer]
            The analyzer to use while simplifying mapped shape expressions.

        Returns
        -------
        result : List[Expr]
            The mapped shape
        """
        return _ffi_api.IndexMapMapShape(self, shape, analyzer)

    def map_tensor(self, arr_src: Tensor) -> Tensor:
        """Apply thie index map to transform the layout of the input Tensor

        Parameters
        ----------
        arr_src : runtime.Tensor
            The Tensor to be transformed

        Returns
        -------
        arr_dst : runtime.Tensor
            The transformed Tensor
        """
        return _ffi_api.IndexMapMapTensor(self, arr_src)

    def inverse(self, shape: list[Range | Expr], analyzer=None) -> "IndexMap":
        """Return the inverse of the map

        Throws an error if the function is not bijective.

        Parameters
        ----------
        shape: List[Union[Range,Expr]]

            The region over which the inverse should be determined.
            Used for validating that the mapping is bijective over
            this range.
        analyzer : Optional[tvm.arith.Analyzer]
            The analyzer to use while deriving and validating the inverse.

        Returns
        -------
        inverse : IndexMap

            The inverse
        """

        shape = [dim if isinstance(dim, Range) else Range(0, dim) for dim in shape]
        return _ffi_api.IndexMapInverse(self, shape, analyzer)

    def non_surjective_inverse(
        self, shape: list[Range | Expr], analyzer=None
    ) -> tuple["IndexMap", Expr]:
        """Return the inverse of the map

        Can be applied to transformations that introduce padding.

        Parameters
        ----------
        shape: List[Union[Range,Expr]]

            The region over which the inverse should be determined.
            Used for determining the predicate.
        analyzer : Optional[tvm.arith.Analyzer]
            The analyzer to use while deriving the inverse and padding predicate.

        Returns
        -------
        result : Tuple[IndexMap, Expr]

            The inverse, and a predicate for which the inverse maps to
            a valid index in the input range.

        Examples
        --------

        .. code-block:: python

            index_map = IndexMap.from_func(lambda i: [i//4, i%4])
            inverse_map, predicate = index_map.non_surjective_inverse([14])
            assert inverse_map.is_equivalent_to(IndexMap.from_func(lambda j,k: [4*j + k])
            print(predicate) # Prints "(axis0==3) && (axis2 >= 2)"
        """

        shape = [dim if isinstance(dim, Range) else Range(0, dim) for dim in shape]
        return _ffi_api.IndexMapNonSurjectiveInverse(self, shape, analyzer)
