# Copyright (c) OpenMMLab. All rights reserved.
from copy import deepcopy
from functools import partial
from typing import Any, Dict, Optional, Sequence, Tuple, Union

import torch

from juxtapose.mmdeploy.core import RewriterContext, patch_model
from juxtapose.mmdeploy.utils import IR, Backend, get_ir_config, get_root_logger
from ..core import PIPELINE_MANAGER


@PIPELINE_MANAGER.register_pipeline()
def trace(
    func: torch.nn.Module,
    inputs: Union[torch.Tensor, Tuple],
    output_path_prefix: Optional[str] = None,
    backend: Union[Backend, str] = "default",
    input_metas: Optional[Dict] = None,
    context_info: Dict = dict(),
    check_trace: bool = True,
    check_tolerance: float = 1e-05,
) -> torch.jit.TracedModule:
    """A wrapper of `torch.jit.trace` with some enhancement.

    Examples:
        >>> from juxtapose.mmdeploy.apis.torch_jit import trace
        >>>
        >>> func = create_model()
        >>> inputs = get_input_tensor()
        >>>
        >>> jit_model = trace(
        >>>     func,
        >>>     inputs,
        >>>     backend='torchscript',
        >>>     check_trace=False)
        >>>

    Args:
        func (torch.nn.Module): A Python function or `torch.nn.Module` that
            will be run with `example_inputs`.
        inputs (torch.Tensor, Tuple): A tuple of example inputs that will be
            passed to the function while tracing.
        output_path_prefix (str): The model would be serialized in
            `<output_path_prefix>.pth`, None if you don't want to
            save the model.
        backend (Backend|str): Which backend will the graph be used. Different
            backend would generate different graph.
        input_metas (Dict): The constant inputs of the model.
        context_info (Dict): The information that would be used in the context
            of exporting.
        check_trace (bool): Check if the same inputs run through traced code
            produce the same outputs.
        check_tolerance (float): Floating-point comparison tolerance to use in
            the checker procedure.

    Returns:
        torch.jit.TracedModule: The traced torch jit model.
    """
    logger = get_root_logger()
    logger.info("Export PyTorch model to torchscript.")

    def _add_or_update(cfg: dict, key: str, val: Any):
        if key in cfg and isinstance(cfg[key], dict) and isinstance(val, dict):
            cfg[key].update(val)
        else:
            cfg[key] = val

    context_info = deepcopy(context_info)
    deploy_cfg = context_info.pop("deploy_cfg", dict())
    ir_config = dict(type="torchscript")
    _add_or_update(deploy_cfg, "ir_config", ir_config)

    if isinstance(backend, Backend):
        backend = backend.value
    backend_config = dict(type=backend)
    _add_or_update(deploy_cfg, "backend_config", backend_config)

    context_info["cfg"] = deploy_cfg
    if "backend" not in context_info:
        context_info["backend"] = backend
    elif context_info["backend"] != backend:
        logger.warning(
            f'Find backend {context_info["backend"]} in context_info.'
            f" Expect {backend}."
        )
    if "ir" not in context_info:
        context_info["ir"] = IR.TORCHSCRIPT
    elif context_info["ir"] != backend:
        logger.warning(
            f'Find ir {context_info["ir"]} in context_info.'
            f" Expect {IR.TORCHSCRIPT}."
        )

    # patch model
    if isinstance(func, torch.nn.Module):
        ir = IR.get(get_ir_config(deploy_cfg)["type"])
        func = patch_model(func, cfg=deploy_cfg, backend=backend, ir=ir)

    with RewriterContext(**context_info), torch.no_grad():
        # patch input_metas
        if input_metas is not None:
            assert isinstance(
                input_metas, dict
            ), f"Expect input_metas type is dict, get {type(input_metas)}."
            model_forward = func.forward
            func.forward = partial(func.forward, **input_metas)

        # for exporting models with weight that depends on inputs
        func(*inputs) if isinstance(inputs, Sequence) else func(inputs)
        ts_model = torch.jit.trace(
            func, inputs, check_trace=check_trace, check_tolerance=check_tolerance
        )

        if input_metas is not None:
            func.forward = model_forward

    # save model
    if output_path_prefix is not None:
        output_path = output_path_prefix + ".pt"
        logger.info(f"Save PyTorch model: {output_path}.")
        torch.jit.save(ts_model, output_path)

    return ts_model
