# Copyright (C) 2022 Intel Corporation
# SPDX-License-Identifier: Apache-2.0
#

import os
import time
from collections.abc import Mapping
from copy import deepcopy
from functools import partial
from subprocess import CalledProcessError
from typing import Any, Callable, Dict, List, Optional, Tuple, Union

import mmcv
import numpy as np
import onnx
import torch
from mmcv.parallel import collate, scatter

from .utils import numpy_2_list
from .utils.mmdeploy import (
    is_mmdeploy_enabled,
    mmdeploy_init_model_helper,
    update_deploy_cfg,
)
from .utils.onnx import prepare_onnx_for_openvino


class NaiveExporter:
    @staticmethod
    def export2openvino(
        output_dir: str,
        model_builder: Callable,
        cfg: mmcv.Config,
        input_data: Dict[Any, Any],
        *,
        precision: str = "FP32",
        model_name: str = "model",
        input_names: Optional[List[str]] = None,
        output_names: Optional[List[str]] = None,
        opset_version: int = 11,
        dynamic_axes: Dict[Any, Any] = {},
        mo_transforms: str = "",
    ):
        input_data = scatter(collate([input_data], samples_per_gpu=1), [-1])[0]

        model = model_builder(cfg)
        model = model.cpu().eval()

        onnx_path = NaiveExporter.torch2onnx(
            output_dir,
            model,
            input_data,
            model_name=model_name,
            input_names=input_names,
            output_names=output_names,
            opset_version=opset_version,
            dynamic_axes=dynamic_axes,
        )

        def get_normalize_cfg(cfg):
            def _get_normalize_cfg(cfg_):
                out = []
                if isinstance(cfg_, Mapping):
                    for key, value in cfg_.items():
                        if isinstance(value, (Mapping, list)):
                            out += _get_normalize_cfg(value)
                        if key == "type" and value == "Normalize":
                            return [cfg_]
                elif isinstance(cfg_, list):
                    for value in cfg_:
                        if isinstance(value, (Mapping, list)):
                            out += _get_normalize_cfg(value)
                return out

            cfg = _get_normalize_cfg(cfg)
            assert len(cfg) == 1
            return cfg[0]

        mo_args = {}

        normalize_cfg = get_normalize_cfg(cfg.data.test)
        if normalize_cfg.get("mean", None) is not None:
            mo_args["mean_values"] = normalize_cfg.get("mean")
        if normalize_cfg.get("std", None) is not None:
            mo_args["scale_values"] = normalize_cfg.get("std")
        if normalize_cfg.get("to_rgb", False):
            mo_args["reverse_input_channels"] = None

        if precision == "FP16":
            mo_args["compress_to_fp16"] = None
        if mo_transforms:
            mo_args["transform"] = mo_transforms

        NaiveExporter.onnx2openvino(
            output_dir,
            onnx_path,
            model_name=model_name,
            **mo_args,
        )

    @staticmethod
    def torch2onnx(
        output_dir: str,
        model: torch.nn.Module,
        input_data: Dict[Any, Any],
        *,
        model_name: str = "model",
        input_names: Optional[List[str]] = None,
        output_names: Optional[List[str]] = None,
        opset_version: int = 11,
        dynamic_axes: Dict[Any, Any] = {},
        verbose: bool = False,
        **onnx_options,
    ) -> str:

        img_metas = input_data.get("img_metas")
        numpy_2_list(img_metas)
        imgs = input_data.get("img")
        model.forward = partial(model.forward, img_metas=img_metas, return_loss=False)

        onnx_file_name = model_name + ".onnx"
        torch.onnx.export(
            model,
            imgs,
            os.path.join(output_dir, onnx_file_name),
            verbose=verbose,
            export_params=True,
            input_names=input_names,
            output_names=output_names,
            dynamic_axes=dynamic_axes,
            opset_version=opset_version,
            operator_export_type=torch.onnx.OperatorExportTypes.ONNX,
            **onnx_options,
        )
        model.__dict__.pop("forward")
        return os.path.join(output_dir, onnx_file_name)

    @staticmethod
    def onnx2openvino(
        output_dir: str,
        onnx_path: str,
        *,
        model_name: str = "model",
        **openvino_options,
    ) -> Tuple[str, str]:
        from otx.mpa.utils import mo_wrapper

        mo_args = {
            "input_model": onnx_path,
            "output_dir": output_dir,
            "model_name": model_name,
        }
        mo_args.update(openvino_options)

        ret, msg = mo_wrapper.generate_ir(output_dir, output_dir, silent=False, **mo_args)
        if ret != 0:
            raise ValueError(msg)
        return (
            os.path.join(output_dir, model_name + ".xml"),
            os.path.join(output_dir, model_name + ".bin"),
        )


if is_mmdeploy_enabled():
    import mmdeploy.apis.openvino as openvino_api
    from mmdeploy.apis import (
        build_task_processor,
        extract_model,
        get_predefined_partition_cfg,
        torch2onnx,
    )
    from mmdeploy.apis.openvino import get_input_info_from_cfg, get_mo_options_from_cfg
    from mmdeploy.core import FUNCTION_REWRITER
    from mmdeploy.utils import get_backend_config, get_ir_config, get_partition_config

    class MMdeployExporter:
        @staticmethod
        def export2openvino(
            output_dir: str,
            model_builder: Callable,
            cfg: mmcv.Config,
            deploy_cfg: mmcv.Config,
            *,
            model_name: str = "model",
        ):

            task_processor = build_task_processor(cfg, deploy_cfg, "cpu")

            def helper(*args, **kwargs):
                return mmdeploy_init_model_helper(*args, **kwargs, model_builder=model_builder)

            task_processor.__class__.init_pytorch_model = helper

            input_data_cfg = deploy_cfg.pop(
                "input_data",
                {"shape": (128, 128, 3), "file_path": None},
            )

            if input_data_cfg.get("file_path"):
                import cv2

                input_data = cv2.imread(input_data_cfg.get("file_path"))
                # image assumed to be RGB format under OTX
                input_data = cv2.cvtColor(input_data, cv2.COLOR_BGR2RGB)
            else:
                input_data = np.zeros(input_data_cfg["shape"], dtype=np.uint8)

            onnx_paths = []
            onnx_paths.append(
                MMdeployExporter.torch2onnx(
                    output_dir,
                    input_data,
                    cfg,
                    deploy_cfg,
                    model_name=model_name,
                )
            )

            partition_cfgs = get_partition_config(deploy_cfg)
            if partition_cfgs:
                partition_cfgs = partition_cfgs.get("partition_cfg", None)
                onnx_paths.extend(
                    MMdeployExporter.partition_onnx(
                        output_dir,
                        onnx_paths[0],
                        partition_cfgs,
                    )
                )

            for i, onnx_path in enumerate(onnx_paths):
                mo_options = {}
                if i > 0:
                    mo_options = partition_cfgs[i - 1].get("mo_options", {})
                deploy_cfg_ = deepcopy(deploy_cfg)
                update_deploy_cfg(onnx_path, deploy_cfg_, mo_options)
                MMdeployExporter.onnx2openvino(
                    output_dir,
                    onnx_path,
                    deploy_cfg_,
                )

        @staticmethod
        def torch2onnx(
            output_dir: str,
            input_data: Any,
            cfg: mmcv.Config,
            deploy_cfg: mmcv.Config,
            *,
            model_name: str = "model",
        ) -> str:
            onnx_file_name = model_name + ".onnx"
            torch2onnx(
                input_data,
                output_dir,
                onnx_file_name,
                deploy_cfg=deploy_cfg,
                model_cfg=cfg,
                model_checkpoint=cfg.get("load_from", None),
                device="cpu",
            )
            return os.path.join(output_dir, onnx_file_name)

        @staticmethod
        def partition_onnx(
            output_dir,
            onnx_path: str,
            partition_cfgs: Union[mmcv.ConfigDict, List[mmcv.ConfigDict]],
        ) -> Tuple[str, ...]:
            partitioned_paths = []

            if not isinstance(partition_cfgs, list):
                partition_cfgs = [partition_cfgs]

            for partition_cfg in partition_cfgs:
                save_file = partition_cfg["save_file"]
                save_path = os.path.join(output_dir, save_file)
                start = partition_cfg["start"]
                end = partition_cfg["end"]
                dynamic_axes = partition_cfg.get("dynamic_axes", None)

                extract_model(onnx_path, start, end, dynamic_axes=dynamic_axes, save_file=save_path)
                partitioned_paths.append(save_path)
            return tuple(partitioned_paths)

        @staticmethod
        def onnx2openvino(
            output_dir: str,
            onnx_path: str,
            deploy_cfg: Union[str, mmcv.Config],
            *,
            model_name: Optional[str] = None,
        ) -> Tuple[str, str]:

            input_info = get_input_info_from_cfg(deploy_cfg)
            output_names = get_ir_config(deploy_cfg).output_names
            mo_options = get_mo_options_from_cfg(deploy_cfg)

            if not model_name:
                model_name = os.path.basename(onnx_path).replace(".onnx", "")
            mo_options.args += f'--model_name "{model_name}" '

            onnx_ready_path = os.path.join(os.path.dirname(onnx_path), f"{model_name}_ready.onnx")
            prepare_onnx_for_openvino(onnx_path, os.path.join(os.path.dirname(onnx_path), f"{model_name}_ready.onnx"))

            try:
                openvino_api.from_onnx(onnx_ready_path, output_dir, input_info, output_names, mo_options)
            except CalledProcessError as e:
                # NOTE: mo returns non zero return code (245) even though it successfully generate IR
                cur_time = time.time()
                time_threshold = 5
                if not (
                    e.returncode == 245
                    and not {model_name + ".bin", model_name + ".xml"} - set(os.listdir(output_dir))
                    and (
                        os.path.getmtime(os.path.join(output_dir, model_name + ".bin")) - cur_time < time_threshold
                        and os.path.getmtime(os.path.join(output_dir, model_name + ".xml")) - cur_time < time_threshold
                    )
                ):
                    raise e

            return (
                os.path.join(output_dir, model_name + ".xml"),
                os.path.join(output_dir, model_name + ".bin"),
            )
