# Copyright (c) OpenMMLab. All rights reserved.
from typing import Any, Optional, Sequence, Tuple, Union

import mmengine
import torch
from torch import Tensor

from juxtapose.mmdeploy.core import FUNCTION_REWRITER
from juxtapose.mmdeploy.core.rewriters.rewriter_utils import LibVersionChecker
from juxtapose.mmdeploy.utils import Backend, load_config


def get_post_processing_params(deploy_cfg: Union[str, mmengine.Config]):
    """Get mmdet post-processing parameters from config.

    Args:
        deploy_cfg (str | mmengine.Config): The path or content of config.

    Returns:
        dict: A dict of parameters for mmdet.
    """
    deploy_cfg = load_config(deploy_cfg)[0]
    codebase_key = "codebase_config"
    assert codebase_key in deploy_cfg
    codebase_config = deploy_cfg[codebase_key]
    post_params = codebase_config.get("post_processing", None)
    assert post_params is not None, "Failed to get `post_processing`."
    return post_params


def clip_bboxes(
    x1: Tensor,
    y1: Tensor,
    x2: Tensor,
    y2: Tensor,
    max_shape: Union[Tensor, Sequence[int]],
):
    """Clip bboxes for onnx.

    Since torch.clamp cannot have dynamic `min` and `max`, we scale the
      boxes by 1/max_shape and clamp in the range [0, 1] if necessary.

    Args:
        x1 (Tensor): The x1 for bounding boxes.
        y1 (Tensor): The y1 for bounding boxes.
        x2 (Tensor): The x2 for bounding boxes.
        y2 (Tensor): The y2 for bounding boxes.
        max_shape (Tensor | Sequence[int]): The (H,W) of original image.

    Returns:
        tuple(Tensor): The clipped x1, y1, x2, y2.
    """
    assert len(max_shape) == 2, "`max_shape` should be [h, w]"
    if isinstance(max_shape, torch.Tensor):
        # scale by 1/max_shape
        x1 = x1 / max_shape[1]
        y1 = y1 / max_shape[0]
        x2 = x2 / max_shape[1]
        y2 = y2 / max_shape[0]

        # clamp [0, 1]
        x1 = torch.clamp(x1, 0, 1)
        y1 = torch.clamp(y1, 0, 1)
        x2 = torch.clamp(x2, 0, 1)
        y2 = torch.clamp(y2, 0, 1)

        # scale back
        x1 = x1 * max_shape[1]
        y1 = y1 * max_shape[0]
        x2 = x2 * max_shape[1]
        y2 = y2 * max_shape[0]
    else:
        x1 = torch.clamp(x1, 0, max_shape[1])
        y1 = torch.clamp(y1, 0, max_shape[0])
        x2 = torch.clamp(x2, 0, max_shape[1])
        y2 = torch.clamp(y2, 0, max_shape[0])
    return x1, y1, x2, y2


@FUNCTION_REWRITER.register_rewriter(
    func_name="mmdeploy.codebase.mmdet.deploy.utils.clip_bboxes",
    backend="tensorrt",
    extra_checkers=LibVersionChecker("tensorrt", min_version="8.2"),
)
def clip_bboxes__trt8(
    x1: Tensor,
    y1: Tensor,
    x2: Tensor,
    y2: Tensor,
    max_shape: Union[Tensor, Sequence[int]],
):
    """Clip bboxes for onnx. From TensorRT 8 we can do the operators on the
    tensors directly.

    Args:
        ctx (ContextCaller): The context with additional information.
        x1 (Tensor): The x1 for bounding boxes.
        y1 (Tensor): The y1 for bounding boxes.
        x2 (Tensor): The x2 for bounding boxes.
        y2 (Tensor): The y2 for bounding boxes.
        max_shape (Tensor | Sequence[int]): The (H,W) of original image.

    Returns:
        tuple(Tensor): The clipped x1, y1, x2, y2.
    """
    assert len(max_shape) == 2, "`max_shape` should be [h, w]"
    x1 = torch.clamp(x1, 0, max_shape[1])
    y1 = torch.clamp(y1, 0, max_shape[0])
    x2 = torch.clamp(x2, 0, max_shape[1])
    y2 = torch.clamp(y2, 0, max_shape[0])
    return x1, y1, x2, y2


def pad_with_value(
    x: Tensor, pad_dim: int, pad_size: int, pad_value: Optional[Any] = None
):
    """Pad a tensor with a value along some dim.

    Args:
        x (Tensor): Input tensor.
        pad_dim (int): Along which dim to pad.
        pad_size (int): To which size to pad.
        pad_value (Any): Filled value for padding. Defaults to `None`.

    Returns:
        Tensor: Padded tensor.
    """
    x_shape = list(x.shape)
    pad_shape = x_shape[:pad_dim] + [pad_size] + x_shape[pad_dim + 1 :]
    x_pad = x.new_zeros(pad_shape)
    if pad_value is not None:
        x_pad = x_pad + pad_value
    x = torch.cat([x, x_pad], dim=pad_dim)
    return x


def pad_with_value_if_necessary(
    x: Tensor, pad_dim: int, pad_size: int, pad_value: Optional[Any] = None
) -> Tensor:
    """Pad a tensor with a value along some dim if necessary.

    Args:
        x (Tensor): Input tensor.
        pad_dim (int): Along which dim to pad.
        pad_size (int): To which size to pad.
        pad_value (Any): Filled value for padding. Defaults to `None`.

    Returns:
        Tensor: Padded tensor.
    """
    return __pad_with_value_if_necessary(
        x, pad_dim, pad_size=pad_size, pad_value=pad_value
    )


def __pad_with_value_if_necessary(
    x: Tensor, pad_dim: int, pad_size: int, pad_value: Optional[Any] = None
) -> Tensor:
    """Pad a tensor with a value along some dim, do nothing on default.

    Args:
        x (Tensor): Input tensor.
        pad_dim (int): Along which dim to pad.
        pad_size (int): To which size to pad.
        pad_value (Any): Filled value for padding. Defaults to `None`.

    Returns:
        Tensor: Padded tensor.
    """
    return x


@FUNCTION_REWRITER.register_rewriter(
    "mmdeploy.codebase.mmdet.deploy.utils.__pad_with_value_if_necessary",
    backend=Backend.TENSORRT.value,
)
def __pad_with_value_if_necessary__tensorrt(
    x: Tensor, pad_dim: int, pad_size: int, pad_value: Optional[Any] = None
) -> Tensor:
    """Pad a tensor with a value along some dim.

    Args:
        x (Tensor): Input tensor.
        pad_dim (int): Along which dim to pad.
        pad_size (int): To which size to pad.
        pad_value (Any): Filled value for padding. Defaults to `None`.

    Returns:
        Tensor: Padded tensor.
    """
    return pad_with_value(x, pad_dim, pad_size=pad_size, pad_value=pad_value)


def __gather_topk(
    *inputs: Sequence[torch.Tensor],
    inds: torch.Tensor,
    batch_size: int,
    is_batched: bool = True,
) -> Tuple[torch.Tensor]:
    """The default implementation of gather_topk."""
    if is_batched:
        batch_inds = torch.arange(batch_size, device=inds.device).unsqueeze(-1)
        outputs = [x[batch_inds, inds, ...] if x is not None else None for x in inputs]
    else:
        prior_inds = inds.new_zeros((1, 1))
        outputs = [x[prior_inds, inds, ...] if x is not None else None for x in inputs]

    return outputs


class TRTGatherTopk(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x: torch.Tensor, inds: torch.Tensor):
        """Implement of gather topk."""
        batch_size = x.size(0)
        batch_inds = torch.arange(batch_size, device=inds.device).unsqueeze(-1)
        return x[batch_inds, inds, ...]

    @staticmethod
    def symbolic(g, x, inds):
        """symbolic of gather topk."""
        out = g.op("mmdeploy::GatherTopk", x, inds, outputs=1)

        return out


@FUNCTION_REWRITER.register_rewriter(
    "mmdeploy.codebase.mmdet.deploy.utils.__gather_topk", backend=Backend.TENSORRT.value
)
def __gather_topk__trt(
    *inputs: Sequence[torch.Tensor],
    inds: torch.Tensor,
    batch_size: int,
    is_batched: bool = True,
) -> Tuple[torch.Tensor]:
    """TensorRT gather_topk."""
    ctx = FUNCTION_REWRITER.get_context()
    _ = ctx
    if is_batched:
        index_shape = inds.shape
        index_dim = inds.dim()
        outputs = [None for _ in inputs]
        for i, x in enumerate(inputs):
            if x is None:
                continue
            out = TRTGatherTopk.apply(x, inds).to(x.dtype)
            out_shape = [*index_shape, *x.shape[index_dim:]]
            out = out.reshape(out_shape)
            outputs[i] = out
    else:
        prior_inds = inds.new_zeros((1, 1))
        outputs = [x[prior_inds, inds, ...] if x is not None else None for x in inputs]

    return outputs


@FUNCTION_REWRITER.register_rewriter(
    "mmdeploy.codebase.mmdet.deploy.utils.__gather_topk", backend=Backend.ASCEND.value
)
@FUNCTION_REWRITER.register_rewriter(
    "mmdeploy.codebase.mmdet.deploy.utils.__gather_topk", backend=Backend.COREML.value
)
def __gather_topk__nonbatch(
    *inputs: Sequence[torch.Tensor],
    inds: torch.Tensor,
    batch_size: int,
    is_batched: bool = True,
) -> Tuple[torch.Tensor]:
    """Single batch gather_topk."""
    assert batch_size == 1
    inds = inds.squeeze(0)
    outputs = [x[:, inds, ...] if x is not None else None for x in inputs]

    return outputs


def gather_topk(
    *inputs: Sequence[torch.Tensor],
    inds: torch.Tensor,
    batch_size: int,
    is_batched: bool = True,
) -> Tuple[torch.Tensor]:
    """Gather topk of each tensor.

    Args:
        inputs (Sequence[torch.Tensor]): Tensors to be gathered.
        inds (torch.Tensor): Topk index.
        batch_size (int): batch_size.
        is_batched (bool): Inputs is batched or not.

    Returns:
        Tuple[torch.Tensor]: Gathered tensors.
    """
    import mmdeploy

    outputs = juxtapose.mmdeploy.codebase.mmdet.deploy.utils.__gather_topk(
        *inputs, inds=inds, batch_size=batch_size, is_batched=is_batched
    )

    if len(outputs) == 1:
        outputs = outputs[0]
    return outputs
