# Copyright (c) OpenMMLab. All rights reserved.
from typing import Dict, Iterable, Optional, Union

import onnx
import onnx.helper
import onnx.utils

from juxtapose.mmdeploy.apis.core import PIPELINE_MANAGER
from juxtapose.mmdeploy.core.optimizers import (
    attribute_to_dict,
    create_extractor,
    get_new_name,
    parse_extractor_io_string,
    remove_identity,
    remove_imports,
    rename_value,
)
from juxtapose.mmdeploy.utils import get_root_logger


@PIPELINE_MANAGER.register_pipeline()
def extract_partition(
    model: Union[str, onnx.ModelProto],
    start_marker: Union[str, Iterable[str]],
    end_marker: Union[str, Iterable[str]],
    start_name_map: Optional[Dict[str, str]] = None,
    end_name_map: Optional[Dict[str, str]] = None,
    dynamic_axes: Optional[Dict[str, Dict[int, str]]] = None,
    save_file: Optional[str] = None,
) -> onnx.ModelProto:
    """Extract partition-model from an ONNX model.

    The partition-model is defined by the names of the input and output tensors
    exactly.

    Examples:
        >>> from juxtapose.mmdeploy.apis import extract_model
        >>> model = 'work_dir/fastrcnn.onnx'
        >>> start_marker = 'detector:input'
        >>> end_marker = ['extract_feat:output', 'multiclass_nms[0]:input']
        >>> dynamic_axes = {
            'input': {
                0: 'batch',
                2: 'height',
                3: 'width'
            },
            'scores': {
                0: 'batch',
                1: 'num_boxes',
            },
            'boxes': {
                0: 'batch',
                1: 'num_boxes',
            }
        }
        >>> save_file = 'partition_model.onnx'
        >>> extract_partition(model, start_marker, end_marker, \
                dynamic_axes=dynamic_axes, \
                save_file=save_file)

    Args:
        model (str | onnx.ModelProto): Input ONNX model to be extracted.
        start_marker (str | Sequence[str]): Start marker(s) to extract.
        end_marker (str | Sequence[str]): End marker(s) to extract.
        start_name_map (Dict[str, str]): A mapping of start names, defaults to
            `None`.
        end_name_map (Dict[str, str]): A mapping of end names, defaults to
            `None`.
        dynamic_axes (Dict[str, Dict[int, str]]): A dictionary to specify
            dynamic axes of input/output, defaults to `None`.
        save_file (str): A file to save the extracted model, defaults to
            `None`.

    Returns:
        onnx.ModelProto: The extracted model.
    """
    if isinstance(model, str):
        model = onnx.load(model)

    num_value_info = len(model.graph.value_info)
    inputs = []
    outputs = []
    logger = get_root_logger()
    if not isinstance(start_marker, (list, tuple)):
        start_marker = [start_marker]
    for s in start_marker:
        start_name, func_id, start_type = parse_extractor_io_string(s)
        for node in model.graph.node:
            if node.op_type == "Mark":
                attr = attribute_to_dict(node.attribute)
                if (
                    attr["func"] == start_name
                    and attr["type"] == start_type
                    and attr["func_id"] == func_id
                ):
                    name = node.input[0]
                    if name not in inputs:
                        new_name = get_new_name(
                            attr, mark_name=s, name_map=start_name_map
                        )
                        rename_value(model, name, new_name)
                        if not any(
                            [
                                v_info.name == new_name
                                for v_info in model.graph.value_info
                            ]
                        ):
                            new_val_info = onnx.helper.make_tensor_value_info(
                                new_name, attr["dtype"], attr["shape"]
                            )
                            model.graph.value_info.append(new_val_info)
                        inputs.append(new_name)

    logger.info(f'inputs: {", ".join(inputs)}')

    # collect outputs
    if not isinstance(end_marker, (list, tuple)):
        end_marker = [end_marker]
    for e in end_marker:
        end_name, func_id, end_type = parse_extractor_io_string(e)
        for node in model.graph.node:
            if node.op_type == "Mark":
                attr = attribute_to_dict(node.attribute)
                if (
                    attr["func"] == end_name
                    and attr["type"] == end_type
                    and attr["func_id"] == func_id
                ):
                    name = node.output[0]
                    if name not in outputs:
                        new_name = get_new_name(
                            attr, mark_name=e, name_map=end_name_map
                        )
                        rename_value(model, name, new_name)
                        if not any(
                            [
                                v_info.name == new_name
                                for v_info in model.graph.value_info
                            ]
                        ):
                            new_val_info = onnx.helper.make_tensor_value_info(
                                new_name, attr["dtype"], attr["shape"]
                            )
                            model.graph.value_info.append(new_val_info)
                        outputs.append(new_name)

    logger.info(f'outputs: {", ".join(outputs)}')

    # replace Mark with Identity
    for node in model.graph.node:
        if node.op_type == "Mark":
            del node.attribute[:]
            node.domain = ""
            node.op_type = "Identity"

    extractor = create_extractor(model)
    extracted_model = extractor.extract_model(inputs, outputs)

    # remove all Identity, this may be done by onnx simplifier
    remove_identity(extracted_model)

    # collect all used inputs
    used = set()
    for node in extracted_model.graph.node:
        for input in node.input:
            used.add(input)

    for output in extracted_model.graph.output:
        used.add(output.name)

    # delete unused inputs
    success = True
    while success:
        success = False
        for i, input in enumerate(extracted_model.graph.input):
            if input.name not in used:
                del extracted_model.graph.input[i]
                success = True
                break

    # eliminate output without shape
    for xs in [extracted_model.graph.output]:
        for x in xs:
            if not x.type.tensor_type.shape.dim:
                logger.info(f"fixing output shape: {x.name}")
                x.CopyFrom(
                    onnx.helper.make_tensor_value_info(
                        x.name, x.type.tensor_type.elem_type, []
                    )
                )

    # eliminate 0-batch dimension, dirty workaround for two-stage detectors
    for input in extracted_model.graph.input:
        if input.name in inputs:
            if input.type.tensor_type.shape.dim[0].dim_value == 0:
                input.type.tensor_type.shape.dim[0].dim_value = 1

    # eliminate duplicated value_info for inputs
    success = True
    # num_value_info == 0 if dynamic shape
    if num_value_info == 0:
        while len(extracted_model.graph.value_info) > 0:
            extracted_model.graph.value_info.pop()
    while success:
        success = False
        for i, x in enumerate(extracted_model.graph.value_info):
            if x.name in inputs:
                del extracted_model.graph.value_info[i]
                success = True
                break

    # dynamic shape support
    if dynamic_axes is not None:
        for input_node in extracted_model.graph.input:
            if input_node.name in dynamic_axes:
                axes = dynamic_axes[input_node.name]
                for k, v in axes.items():
                    input_node.type.tensor_type.shape.dim[k].dim_value = 0
                    input_node.type.tensor_type.shape.dim[k].dim_param = v
        for output_node in extracted_model.graph.output:
            for idx, dim in enumerate(output_node.type.tensor_type.shape.dim):
                dim.dim_value = 0
                dim.dim_param = f"dim_{idx}"

    # remove mmdeploy domain if useless
    remove_imports(extracted_model)

    # save extract_model if save_file is given
    if save_file is not None:
        onnx.save(extracted_model, save_file)

    return extracted_model
