# Copyright (c) OpenMMLab. All rights reserved.
import inspect
from typing import Any, Callable, Dict, Optional, Sequence

import torch

from juxtapose.mmdeploy.core.rewriters import FUNCTION_REWRITER
from juxtapose.mmdeploy.utils import IR, cfg_apply_marks, get_partition_config

MARK_FUNCTION_COUNT = dict()


def reset_mark_function_count():
    """Reset counter of mark function."""
    for k in MARK_FUNCTION_COUNT:
        MARK_FUNCTION_COUNT[k] = 0


TORCH_DTYPE_TO_ONNX = {
    torch.uint8: torch.onnx.TensorProtoDataType.UINT8,
    torch.int8: torch.onnx.TensorProtoDataType.INT8,
    torch.float64: torch.onnx.TensorProtoDataType.DOUBLE,
    torch.float32: torch.onnx.TensorProtoDataType.FLOAT,
    torch.float16: torch.onnx.TensorProtoDataType.FLOAT16,
    torch.int32: torch.onnx.TensorProtoDataType.INT32,
    torch.int64: torch.onnx.TensorProtoDataType.INT64,
    torch.int16: torch.onnx.TensorProtoDataType.INT16,
    torch.bool: torch.onnx.TensorProtoDataType.BOOL,
    torch.complex64: torch.onnx.TensorProtoDataType.COMPLEX64,
    torch.complex128: torch.onnx.TensorProtoDataType.COMPLEX128,
}


class Mark(torch.autograd.Function):
    """The mark node function.

    The function does nothing but inserts a mark node to ONNX model. The mark
    can ease the process of model partition.
    """

    @staticmethod
    def symbolic(g, x, dtype, shape, func, func_id, type, name, id, attrs):
        """Symbolic function for mmdeploy::Mark op."""
        n = g.op(
            "mmdeploy::Mark",
            x,
            dtype_i=TORCH_DTYPE_TO_ONNX[dtype],
            shape_i=shape,
            func_s=func,
            func_id_i=func_id,
            type_s=type,
            name_s=name,
            id_i=id,
            **attrs,
        )
        return n

    @staticmethod
    def forward(ctx, x, *args) -> torch.Tensor:
        """Run forward."""
        return x


@FUNCTION_REWRITER.register_rewriter(
    "mmdeploy.core.optimizers.function_marker.Mark.symbolic"
)
def mark_symbolic(g, x, *args):
    """Rewrite symbolic of mark op."""
    ctx = FUNCTION_REWRITER.get_context()
    if cfg_apply_marks(ctx.cfg):
        return ctx.origin_func(g, x, *args)
    return x


@FUNCTION_REWRITER.register_rewriter(
    "mmdeploy.core.optimizers.function_marker.Mark.forward"
)
def forward_of_mark(
    ctx, x, dtype, shape, func, func_id, type, name, id, attrs
) -> torch.Tensor:
    """Rewrite forward of mark op."""
    rewriter = FUNCTION_REWRITER.get_context()
    deploy_cfg = rewriter.cfg
    # save calib data
    apply_marks = cfg_apply_marks(deploy_cfg)
    create_calib = getattr(rewriter, "create_calib", False)
    if apply_marks and create_calib:
        partition_params = get_partition_config(deploy_cfg)
        assert partition_params is not None, "No partition config."
        partition_type = partition_params["type"]

        from juxtapose.mmdeploy.apis import get_predefined_partition_cfg

        partition_cfgs = get_predefined_partition_cfg(deploy_cfg, partition_type)
        assert hasattr(rewriter, "calib_file")

        for partition_id, partition_cfg in enumerate(partition_cfgs):
            start = partition_cfg["start"]
            if (f"{func}:{type}" not in start) and (
                f"{func}[{func_id}]:{type}" not in start
            ):
                continue

            input_name = name
            dynamic_axes = partition_cfg.get("dynamic_axes", None)
            if dynamic_axes is not None:
                input_name = name
            calib_file = rewriter.calib_file

            calib_data_group = calib_file["calib_data"]
            partition_name = f"partition{partition_id}"

            if partition_name not in calib_data_group:
                calib_data_group.create_group(partition_name)
            partition_group = calib_data_group[partition_name]

            if input_name not in partition_group:
                partition_group.create_group(input_name)
            input_data_group = partition_group[input_name]

            data_id = rewriter.data_id
            x_np = x.detach().cpu().numpy()
            input_data_group.create_dataset(
                str(data_id),
                shape=x_np.shape,
                compression="gzip",
                compression_opts=4,
                data=x_np,
            )

    return rewriter.origin_func(
        ctx, x, dtype, shape, func, func_id, type, name, id, attrs
    )


def mark_tensors(
    xs: Any,
    func: str,
    func_id: int,
    io_type: str,
    ctx: Any,
    attrs: Dict,
    is_inspecting: bool,
    level: int,
) -> tuple:
    """Add mark node recursively.

    Args:
        xs (Any): Input structure which contains tensor.
        func (str): Function name of the function which xs comes from.
        func_id (int): Function index of `func` in the model.
        io_type (str): The io type of xs, `input` or `output`.
        ctx (Any): The context instance.
        attrs (Dict): The extra attributes provided by mark decorator.
        is_inspecting (bool): The names of xs are inspected or not.
        level (int): The recursive level.

    Returns:
        Any: The same structure as xs, all tensor has been replaced with Mark.
    """
    visit = set()
    index = 0

    def impl(ys, prefix, level):
        nonlocal index
        old_index = index
        ret = ys
        prefix = () if level == 0 else prefix

        if isinstance(ys, torch.Tensor):
            if ys not in visit:
                visit.add(ys)
                root = ctx.names[ctx.index]
                name = ".".join(str(x) for x in (root, *prefix))
                ys_shape = tuple(int(s) for s in ys.shape)
                ret = Mark.apply(
                    ys, ys.dtype, ys_shape, func, func_id, io_type, name, index, attrs
                )
                index += 1
        elif isinstance(ys, list):
            ret = [impl(y, prefix + (i,), level + 1) for i, y in enumerate(ys)]
        elif isinstance(ys, tuple):
            ret = tuple(impl(y, prefix + (i,), level + 1) for i, y in enumerate(ys))
        elif isinstance(ys, dict):
            ret = {k: impl(v, prefix + (k,), level + 1) for k, v in ys.items()}

        if level == 0 and (is_inspecting or old_index != index):
            ctx.index += 1

        return ret

    return impl(xs, (), level)


@FUNCTION_REWRITER.register_rewriter(
    "mmdeploy.core.optimizers.function_marker.mark_tensors", ir=IR.TORCHSCRIPT
)
def remove_mark__torchscript(xs: Any, *args, **kwargs):
    """Disable all marks for TorchScript backend.

    As the Node `mark` is not able to be traced, we just return original input
    for the function `mark_tensors`.

    Args:
        xs (Any): Input structure which contains tensor.
    """
    return xs


def mark(
    func_name: Optional[str] = None,
    inputs: Optional[Sequence[str]] = None,
    outputs: Optional[Sequence[str]] = None,
    **attrs,
) -> Callable:
    """The decorator used to add mark node.

    Mark node can be used to support model partition.

    Args:
        func_name (str): The name of the function where marks come from.
        inputs (Sequence[str]): The input names of the marks. The final name \
            might have suffix if inputs is list or dictionary.
        outputs (Sequence[str]): The output names of the marks. The final \
            name might have suffix if outputs is list or dictionary.

    Returns:
        Callable: The process of mark decorator.

    Examples:
        >>> from juxtapose.mmdeploy.core import FUNCTION_REWRITER, mark
        >>> @FUNCTION_REWRITER.register_rewriter(
        >>>     func_name='mmdet.models.roi_heads.ConvFCBBoxHead.forward')
        >>> def forward_of_bbox_head(self, x):
        >>>     ctx = FUNCTION_REWRITER.get_context()
        >>>     @mark(
        >>>         'bbox_head_forward',
        >>>         inputs=['bbox_feats'],
        >>>         outputs=['cls_score', 'bbox_pred'])
        >>>     def _impl():
        >>>         return ctx.origin_func(self, x)
        >>>     return _impl()
    """
    MARK_FUNCTION_COUNT[func_name] = 0

    class Context:
        def __init__(self, names):
            self.names = names
            self.index = 0

    def decorator(f):
        func = func_name if func_name else f.__name__
        is_inspect = False
        if not inputs:
            input_names = list(inspect.signature(f).parameters.keys())
            is_inspect = True
        else:
            input_names = inputs
        output_names = outputs if outputs else func

        # args and retvals match corresponding names at level 0
        args_level, rets_level = -1, -1

        if isinstance(input_names, str):
            input_names = (input_names,)

        if isinstance(output_names, str):
            output_names = (output_names,)
            rets_level += 1

        def g(*args, **kwargs):
            func_id = MARK_FUNCTION_COUNT[func_name]
            MARK_FUNCTION_COUNT[func_name] += 1
            ctx = Context(input_names)
            args = mark_tensors(
                args, func, func_id, "input", ctx, attrs, is_inspect, args_level
            )

            rets = f(*args, **kwargs)

            ctx = Context(output_names)
            func_ret = mark_tensors(
                rets, func, func_id, "output", ctx, attrs, False, rets_level
            )
            return func_ret

        return g

    return decorator
