# 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.
"""Operation class for computation declaration."""

import inspect

# pylint: disable=invalid-name
from numbers import Integral as _Integral

from tvm_ffi import Array

import tvm.arith._ffi_api
import tvm.tirx
import tvm.tirx._ffi_api
from tvm.ir import is_prim_expr
from tvm.runtime import convert

from . import _ffi_api
from . import tag as _tag
from . import tensor as _tensor


def placeholder(shape, dtype=None, name="placeholder"):
    """Construct an empty tensor object.

    Parameters
    ----------
    shape: Tuple of Expr
        The shape of the tensor

    dtype: str, optional
        The data type of the tensor

    name: str, optional
        The name hint of the tensor

    Returns
    -------
    tensor: Tensor
        The created tensor
    """
    dtype = "float32" if dtype is None else dtype
    return _ffi_api.Placeholder(shape, dtype, name)


def compute(shape, fcompute, name="compute", tag="", attrs=None, varargs_names=None):
    """Construct a new tensor by computing over the shape domain.

    The compute rule is result[axis] = fcompute(axis)

    Parameters
    ----------
    shape: Tuple of Expr
        The shape of the tensor

    fcompute: lambda function of indices-> value
        Specifies the input source expression

    name: str, optional
        The name hint of the tensor

    tag: str, optional
        Additional tag information about the compute.

    attrs: dict, optional
        The additional auxiliary attributes about the compute.

    varargs_names: list, optional
        The names to use for each of the varargs. If not supplied, the varargs
        will be called i1, i2, ...

    Returns
    -------
    tensor: Tensor
        The created tensor
    """
    if _tag.TagScope.get_current() is not None:
        if tag != "":
            raise ValueError("nested tag is not allowed for now")
        tag = _tag.TagScope.get_current().tag
    shape = (shape,) if tvm.ir.is_prim_expr(shape) else shape
    # for python3
    shape = tuple([int(s) if isinstance(s, float) else s for s in shape])
    out_ndim = len(shape)

    argspec = inspect.getfullargspec(fcompute)
    if len(argspec.args) == 0 and argspec.varargs is None:
        arg_names = [f"i{i}" for i in range(out_ndim)]
    elif argspec.varargs is not None:
        # if there is a varargs, it takes the remaining dimensions of out_ndim
        num_remaining_args = out_ndim - len(argspec.args)
        if varargs_names is not None:
            if len(varargs_names) != num_remaining_args:
                raise RuntimeError(
                    f"Number of varargs ({num_remaining_args}) does not match number"
                    f"of varargs_names ({len(varargs_names)})"
                )
            arg_names = argspec.args + varargs_names
        else:
            arg_names = argspec.args + [f"i{i}" for i in range(out_ndim - len(argspec.args))]
    else:
        arg_names = argspec.args
        # if there are fewer args than out dimensions, the remaining dimensions
        # are implicitly broadcast
        out_ndim = len(arg_names)
    assert argspec.varkw is None, "Variable keyword arguments not supported in fcompute"
    assert argspec.defaults is None, "Default arguments not supported in fcompute"
    assert len(argspec.kwonlyargs) == 0, "Keyword arguments are not supported in fcompute"

    if out_ndim != len(arg_names):
        raise ValueError(
            "Number of args to fcompute does not match dimension, "
            f"args={len(arg_names)}, dimension={out_ndim}"
        )

    dim_var = [tvm.tirx.IterVar((0, s), x, 0) for x, s in zip(arg_names, shape[:out_ndim])]
    body = fcompute(*[v.var for v in dim_var])

    if not isinstance(body, list | tuple):
        body = [body]
    body = convert(body)
    op_node = _ffi_api.ComputeOp(name, tag, attrs, dim_var, body)

    num = op_node.num_outputs
    outputs = tuple(op_node.output(i) for i in range(num))
    return outputs[0] if num == 1 else outputs


def scan(init, update, state_placeholder, inputs=None, name="scan", tag="", attrs=None):
    """Construct new tensors by scanning over axis.

    Parameters
    ----------
    init: Tensor or list of Tensor
        The initial condition of first init.shape[0] timestamps

    update: Tensor or list of Tensor
        The update rule of the scan given by symbolic tensor.

    state_placeholder: Tensor or list of Tensor
        The placeholder variables used by update.

    inputs: Tensor or list of Tensor, optional
        The list of inputs to the scan. This is not required, but can
        be useful for the compiler to detect scan body faster.

    name: str, optional
        The name hint of the tensor

    tag: str, optional
        Additonal tag information about the compute.

    attrs: dict, optional
        The additional auxiliary attributes about the compute.

    Returns
    -------
    tensor: Tensor or list of Tensors
        The created tensor or tuple of tensors contains multiple outputs.

    Example
    -------
    .. code-block:: python

      # The following code is equivalent to numpy.cumsum
      m = te.var("m")
      n = te.var("n")
      X = te.placeholder((m, n), name="X")
      s_state = te.placeholder((m, n))
      s_init = te.compute((1, n), lambda _, i: X[0, i])
      s_update = te.compute((m, n), lambda t, i: s_state[t-1, i] + X[t, i])
      res = tvm.te.scan(s_init, s_update, s_state, X)
    """
    if _tag.TagScope.get_current() is not None:
        if tag != "":
            raise ValueError("nested tag is not allowed for now")
        tag = _tag.TagScope.get_current().tag
    if isinstance(init, _tensor.Tensor):
        init = [init]
    if isinstance(update, _tensor.Tensor):
        update = [update]
    if isinstance(state_placeholder, _tensor.Tensor):
        state_placeholder = [state_placeholder]
    if isinstance(inputs, _tensor.Tensor):
        inputs = [inputs]
    if inputs is None:
        inputs = []
    if len(init) != len(update) or len(init) != len(state_placeholder):
        raise ValueError("init, update, state_placeholder must have same length")
    axis = tvm.tirx.IterVar((init[0].shape[0], update[0].shape[0]), f"{name}.idx", 3)
    op = _ffi_api.ScanOp(name, tag, attrs, axis, init, update, state_placeholder, inputs)
    res = [op.output(i) for i in range(len(update))]
    return res[0] if len(res) == 1 else res


def extern(
    shape,
    inputs,
    fcompute,
    name="extern",
    dtype=None,
    in_buffers=None,
    out_buffers=None,
    tag="",
    attrs=None,
):
    """Compute several tensors via an extern function.

    Parameters
    ----------
    shape: tuple or list of tuples.
        The shape of the outputs.

    inputs: list of Tensor
        The inputs

    fcompute: lambda function of inputs, outputs-> stmt
        Specifies the IR statement to do the computation.
        See the following note for function signature of fcompute

        .. note::
             **Parameters**

             - **ins** (list of :any:`tvm.tirx.Buffer`) - Placeholder for each inputs
             - **outs** (list of :any:`tvm.tirx.Buffer`) - Placeholder for each outputs

             **Returns**

             - **stmt** (:any:`tvm.tirx.Stmt`) - The statement that carries out array computation.

    name: str, optional
        The name hint of the tensor

    dtype: str or list of str, optional
        The data types of outputs,
        by default dtype will be same as inputs.

    in_buffers: tvm.tirx.Buffer or list of tvm.tirx.Buffer, optional
        Input buffers.

    out_buffers: tvm.tirx.Buffer or list of tvm.tirx.Buffer, optional
        Output buffers.


    tag: str, optional
        Additonal tag information about the compute.

    attrs: dict, optional
        The additional auxiliary attributes about the compute.

    Returns
    -------
    tensor: Tensor or list of Tensors
        The created tensor or tuple of tensors contains multiple outputs.

    Example
    -------
    In the code below, C is generated by calling external PackedFunc
    `tvm.contrib.cblas.matmul`

    .. code-block:: python

        A = te.placeholder((n, l), name="A")
        B = te.placeholder((l, m), name="B")
        C = te.extern((n, m), [A, B],
                       lambda ins, outs: tvm.tirx.call_packed(
                          "tvm.contrib.cblas.matmul",
                            ins[0], ins[1], outs[0], 0, 0), name="C")
    """
    if _tag.TagScope.get_current() is not None:
        if tag != "":
            raise ValueError("nested tag is not allowed for now")
        tag = _tag.TagScope.get_current().tag
    shape = (shape,) if is_prim_expr(shape) or isinstance(shape, _Integral) else shape
    if shape == () or is_prim_expr(shape[0]) or isinstance(shape[0], _Integral):
        shape = [shape]
    if in_buffers is not None:
        in_buffers = [in_buffers] if not isinstance(in_buffers, list) else in_buffers
        if len(inputs) != len(in_buffers):
            raise RuntimeError(
                f"Number of inputs and in_buffers mismatch: {len(inputs)} vs {len(in_buffers)}."
            )
    if out_buffers is not None:
        out_buffers = [out_buffers] if not isinstance(out_buffers, list) else out_buffers
        if len(shape) != len(out_buffers):
            raise RuntimeError(
                f"Number of outputs and out_buffers mismatch: {len(shape)} vs {len(out_buffers)}."
            )
    input_placeholders = in_buffers or []
    output_placeholders = out_buffers or []
    types = set()
    for t in inputs:
        if not isinstance(t, _tensor.Tensor):
            raise ValueError("expect inputs to be tensor")
        if in_buffers is None:
            input_placeholders.append(
                tvm.tirx.decl_buffer(
                    t.shape,
                    t.dtype,
                    t.op.name,
                    elem_offset=tvm.tirx.Var("elem_offset", "int32"),
                    layout=None,
                )
            )
        types.add(t.dtype)

    if out_buffers is None:
        if dtype is None:
            if len(types) != 1:
                raise ValueError("Cannot infer output type, please provide dtype argument")
            infered_type = types.pop()
            dtype = [infered_type for _ in shape]
        if isinstance(dtype, str):
            dtype = [dtype]

        for shp, dt in zip(shape, dtype):
            output_placeholders.append(
                tvm.tirx.decl_buffer(
                    shp,
                    dt,
                    name,
                    elem_offset=tvm.tirx.Var("elem_offset", "int32"),
                    layout=None,
                )
            )
    body = fcompute(input_placeholders, output_placeholders)
    if tvm.ir.is_prim_expr(body):
        body = tvm.tirx.Evaluate(body)
    if not isinstance(body, tvm.tirx.Stmt):
        raise ValueError(
            f"Function '{fcompute.__name__}' should return Expr or Stmt, but it returned "
            f"'{type(body)}'"
        )

    op = _ffi_api.ExternOp(name, tag, attrs, inputs, input_placeholders, output_placeholders, body)
    res = [op.output(i) for i in range(len(output_placeholders))]
    return res[0] if len(res) == 1 else res


def extern_primfunc(input_tensors: list[_tensor.Tensor], primfunc: tvm.tirx.PrimFunc, **kwargs):
    """Compute tensors via a schedulable TIR PrimFunc

    Parameters
    ----------
    input_tensors: list of Tensor
        Input tensors that map to the corresponding primfunc input params.

    primfunc: PrimFunc
        The TIR PrimFunc

    Returns
    -------
    tensor: Tensor or list of Tensors
        The created tensor or tuple of tensors if it contains multiple outputs.

    Example
    -------
    In the code below, a TVMScript defined TIR PrimFunc is inlined into
    a TE ExternOp. Applying te.create_prim_func on this

    .. code-block:: python

        A = te.placeholder((128, 128), name="A")
        B = te.placeholder((128, 128), name="B")

        @T.prim_func(s_tir=True)
        def before_split(a: T.handle, b: T.handle) -> None:
            A = T.match_buffer(a, (128, 128))
            B = T.match_buffer(b, (128, 128))
            for i, j in T.grid(128, 128):
                with T.sblock("B"):
                    vi, vj = T.axis.remap("SS", [i, j])
                    B[vi, vj] = A[vi, vj] * 2.0

        C = te.extern_primfunc([A, B], func)
    """

    # dt_access_map and primfunc.buffer_map are unordered, so use order from primfunc.params
    dt_access_map = tvm.arith._ffi_api.DomainTouchedAccessMap(primfunc)
    ordered_buffers = [primfunc.buffer_map[param] for param in primfunc.params]
    in_buffers = [buf for buf in ordered_buffers if len(dt_access_map[buf][0])]
    out_buffers = [buf for buf in ordered_buffers if len(dt_access_map[buf][1])]
    assert in_buffers, "PrimFunc has no input buffers"
    assert out_buffers, "PrimFunc has no output buffers"

    outputs = []
    inplace = []
    input_buffers = in_buffers
    for obuf in out_buffers:
        if obuf in in_buffers:
            inplace.append(obuf)
        else:
            outputs.append(obuf)

    if not outputs:
        iobuf = inplace.pop()
        input_buffers.remove(iobuf)
        outputs = [iobuf]

    assert len(input_buffers) == len(input_tensors), (
        "The number of provided input input_tensors does not match the number of ",
        "input buffers in the primfunc",
    )
    for tensor, buffer in zip(input_tensors, input_buffers):
        # TODO(csullivan): Can a stronger comparison between Tensor<>Buffer be made?
        assert len(tensor.shape) == len(buffer.shape)
        for d1, d2 in zip(tensor.shape, buffer.shape):
            assert d1 == d2, (
                "The input input_tensors provided do not match the input buffers in the ",
                "primfunc. Please check that the order of input te.Input_Tensors and the ",
                "order of the primfunc variables in the params list agree.",
            )
    output = extern(
        [buf.shape for buf in outputs],
        input_tensors,
        lambda ins, outs: primfunc.body,
        in_buffers=input_buffers,
        out_buffers=outputs,
        **kwargs,
    )
    return output


def var(name="tindex", dtype="int32", span=None):
    """Create a new variable with specified name and dtype

    Parameters
    ----------
    name : str
        The name

    dtype : str
        The data type

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

    Returns
    -------
    var : tirx.Var
        The result symbolic variable.
    """
    return tvm.tirx.Var(name, dtype, span)


def const(value, dtype="int32", span=None):
    """Create a new constant with specified value and dtype

    Parameters
    ----------
    value : Union[bool, int, float, numpy.ndarray, tvm.runtime.Tensor]
        The constant value.

    dtype : str
        The data type

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

    Returns
    -------
    const : Expr
        The result constant expr.
    """
    return tvm.tirx.const(value, dtype, span)


def thread_axis(dom=None, tag="", name="", span=None):
    """Create a new IterVar to represent thread index.

    Parameters
    ----------
    dom : Range or str
        The domain of iteration
        When str is passed, dom is set to None and str is used as tag

    tag : str, optional
        The thread tag

    name : str, optional
        The name of the var.

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

    Returns
    -------
    axis : IterVar
        The thread itervar.
    """
    if isinstance(dom, str):
        tag, dom = dom, None
    if not tag:
        raise ValueError("tag must be given as Positional or keyword argument")
    name = name if name else tag
    return tvm.tirx.IterVar(dom, name, 1, tag, span)


def reduce_axis(dom, name="rv", thread_tag="", span=None):
    """Create a new IterVar for reduction.

    Parameters
    ----------
    dom : Range
        The domain of iteration.

    name : str
        The name of the variable.

    thread_tag : Optional[str]
        The name of the thread_tag.

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

    Returns
    -------
    axis : IterVar
        An iteration variable representing the value.
    """
    return tvm.tirx.IterVar(dom, name, 2, thread_tag, span)


def create_prim_func(
    ops: list[_tensor.Tensor | tvm.tirx.Var], index_dtype_override: str | None = None
) -> tvm.tirx.PrimFunc:
    """Create a TensorIR PrimFunc from tensor expression

    Parameters
    ----------
    ops : List[Union[_tensor.Tensor, tvm.tirx.Var]]
        The source expression.

    Example
    -------
    We define a matmul kernel using following code:

    .. code-block:: python

        import tvm
        from tvm import te
        from tvm.te import create_prim_func
        import tvm.script

        A = te.placeholder((128, 128), name="A")
        B = te.placeholder((128, 128), name="B")
        k = te.reduce_axis((0, 128), "k")
        C = te.compute((128, 128), lambda x, y: te.sum(A[x, k] * B[y, k], axis=k), name="C")
        func = create_prim_func([A, B, C])
        print(func.script())

    If we want to use TensorIR schedule to do transformations on such kernel,
    we need to use `create_prim_func([A, B, C])` to create a schedulable PrimFunc.
    The generated function looks like:

    .. code-block:: python

        @T.prim_func(s_tir=True)
        def tir_matmul(a: T.handle, b: T.handle, c: T.handle) -> None:
            A = T.match_buffer(a, (128, 128))
            B = T.match_buffer(b, (128, 128))
            C = T.match_buffer(c, (128, 128))

            for i, j, k in T.grid(128, 128, 128):
                with T.sblock():
                    vi, vj, vk = T.axis.remap("SSR", [i, j, k])
                    with T.init():
                        C[vi, vj] = 0.0
                    C[vi, vj] += A[vi, vk] * B[vj, vk]

    Returns
    -------
    func : tirx.PrimFunc
        The created function.
    """
    if not isinstance(ops, list | tuple | Array):
        ops = [ops]
    return _ffi_api.CreatePrimFunc(ops, index_dtype_override)


AXIS_SEPARATOR = tvm.tirx.IndexMap.AXIS_SEPARATOR