# 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=redefined-builtin
"""Operators for distributed Relax."""

from tvm.ir import Call
from tvm.relax.distributed import DeviceMesh, DTensorType, Placement

from ...expr import Expr, GlobalVar, ShapeExpr
from ...expr import Tuple as RxTuple
from ...utils import convert_to_expr
from . import _ffi_api


def annotate_sharding(input: Expr, device_mesh: DeviceMesh, placement: Placement) -> Expr:
    """Annotate sharding plan for tensor

    Parameters
    ----------
    input : relax.Expr
      The input tensor.
    device_mesh: DeviceMesh
      The device mesh of the sharding plan
    placement: Placement
      The placement of the sharding plan

    Returns
    -------
    result : relax.Expr
      The tensor unmodified.
    """
    return _ffi_api.annotate_sharding(input, device_mesh, placement)  # type: ignore


def redistribute(input: Expr, device_mesh: DeviceMesh, placement: Placement) -> Expr:
    """Redistribute tensor

    Parameters
    ----------
    input : relax.Expr
      The input tensor.
    device_mesh: DeviceMesh
      The device mesh after redistribution
    placement: Placement
      The placement after redistribution
    Returns
    -------
    result : relax.Expr
      The tensor after redistribution.
    """
    return _ffi_api.redistribute(input, device_mesh, placement)  # type: ignore


def call_tir_local_view(
    gvar: GlobalVar,
    args: Expr,
    out_ty: DTensorType | list[DTensorType],
    tir_vars: ShapeExpr | tuple[Expr] | list[Expr] | None = None,
) -> Call:
    """
    Call a tirx.prim_func and return the output. The prim_func should be a worker-local function
    that is actually executed on each worker, instead of the unpartitioned function.
    The output of this operator is DTensor or a tuple of DTensors.

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

    args : Expr
        The input arguments.

    out_ty : Union[DTensorType, List[DTensorType]]
        The type information of the call_tir output.
        It should be a single or a list of DTensorType. 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_local_view operator.
    """
    if isinstance(args, tuple | list):
        args = RxTuple([convert_to_expr(a) for a in args])
    elif isinstance(args, Expr) and not isinstance(args, RxTuple):  # type: ignore
        args = RxTuple((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_local_view(gvar, args, out_ty, tir_vars)  # type: ignore


def redistribute_replica_to_shard(input: Expr, num_workers: int, axis: int) -> Expr:
    """Slice tensor into several parts along one axis,
        and each worker takes one part.
        input.ty.shape[axis] % num_workers == 0 is required.
        Each worker must have an identical copy of the input.
        This is a specialized version of redistribute op.

    Parameters
    ----------
    input : relax.Expr
      The buffer to be sliced into equal parts.

    num_worker : int
      The number of workers, i.e. the number of parts the given buffer should be sliced into.

    axis : int
      The axis of the tensor to be sliced.

    Returns
    -------
    result : relax.Expr
      Sliced Tensor kept by each device.
    """
    return _ffi_api.redistribute_replica_to_shard(input, num_workers, axis)