# Copyright (c) OpenMMLab. All rights reserved.
from typing import Optional

import torch

from juxtapose.mmdeploy.core import FUNCTION_REWRITER
from juxtapose.mmdeploy.utils import get_root_logger


@FUNCTION_REWRITER.register_rewriter(func_name="torch.topk", backend="default")
@FUNCTION_REWRITER.register_rewriter(func_name="torch.Tensor.topk", backend="default")
def topk__dynamic(
    input: torch.Tensor,
    k: int,
    dim: Optional[int] = None,
    largest: bool = True,
    sorted: bool = True,
):
    """Rewrite `topk` for default backend.

    Cast k to tensor and make sure k is smaller than input.shape[dim].
    """
    ctx = FUNCTION_REWRITER.get_context()

    if dim is None:
        dim = int(input.ndim - 1)
    size = input.shape[dim]
    if not isinstance(k, torch.Tensor):
        k = torch.tensor(k, device=input.device, dtype=torch.long)
    # Always keep topk op for dynamic input
    if isinstance(size, torch.Tensor):
        # size would be treated as cpu tensor, trick to avoid that.
        size = k.new_zeros(()) + size
    k = torch.where(k < size, k, size)
    return ctx.origin_func(input, k, dim=dim, largest=largest, sorted=sorted)


@FUNCTION_REWRITER.register_rewriter(func_name="torch.topk", backend="tensorrt")
@FUNCTION_REWRITER.register_rewriter(func_name="torch.Tensor.topk", backend="tensorrt")
def topk__tensorrt(
    input: torch.Tensor,
    k: int,
    dim: Optional[int] = None,
    largest: bool = True,
    sorted: bool = True,
):
    """Rewrite `topk` for TensorRT backend.

    TensorRT does not support topk with dynamic k. This function cast k to
    constant integer.
    """
    ctx = FUNCTION_REWRITER.get_context()
    # https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#topKsetup
    from juxtapose.mmdeploy.utils.constants import TENSORRT_MAX_TOPK

    if dim is None:
        dim = int(input.ndim - 1)
    size = input.shape[dim]
    if k > size:
        k = size
    if not isinstance(k, int):
        k = int(k)
    if k > TENSORRT_MAX_TOPK:
        logger = get_root_logger()
        logger.warning(
            f"Maximum K of TopK in TensorRT is {TENSORRT_MAX_TOPK}, "
            f"but given {k}. Note that k will be set "
            f"to {TENSORRT_MAX_TOPK}."
        )
        k = TENSORRT_MAX_TOPK

    return ctx.origin_func(input, k, dim=dim, largest=largest, sorted=sorted)


@FUNCTION_REWRITER.register_rewriter(func_name="torch.topk", backend="coreml")
@FUNCTION_REWRITER.register_rewriter(func_name="torch.Tensor.topk", backend="coreml")
def topk__coreml(
    input: torch.Tensor,
    k: int,
    dim: Optional[int] = None,
    largest: bool = True,
    sorted: bool = True,
):
    """Rewrite `topk` for coreml.

    Cast k to tensor and make sure k is smaller than input.shape[dim].
    """
    ctx = FUNCTION_REWRITER.get_context()

    if dim is None:
        dim = int(input.ndim - 1)
    size = input.shape[dim]
    if not isinstance(k, torch.Tensor):
        k = torch.tensor(k, device=input.device, dtype=torch.long)
    # Always keep topk op for dynamic input
    k = torch.where(k < size, k, size)
    return ctx.origin_func(input, k, dim=dim, largest=largest, sorted=sorted)
