# Copyright (c) OpenMMLab. All rights reserved.

import copy
import os
from collections import defaultdict
from typing import Callable, Dict, Optional, Sequence, Tuple, Union

import mmcv
import mmengine
import numpy as np
import torch
from mmengine.model import BaseDataPreprocessor
from mmengine.registry import Registry

from juxtapose.mmdeploy.codebase.base import CODEBASE, BaseTask, MMCodebase
from juxtapose.mmdeploy.utils import Codebase, Task, get_input_shape, get_root_logger


def process_model_config(
    model_cfg: mmengine.Config,
    imgs: Union[Sequence[str], Sequence[np.ndarray]],
    input_shape: Optional[Sequence[int]] = None,
):
    """Process the model config for sdk model.

    Args:
        model_cfg (mmengine.Config): The model config.
        imgs (Sequence[str] | Sequence[np.ndarray]): Input image(s), accepted
            data type are List[str], List[np.ndarray].
        input_shape (list[int]): A list of two integer in (width, height)
            format specifying input shape. Default: None.

    Returns:
        mmengine.Config: the model config after processing.
    """
    cfg = copy.deepcopy(model_cfg)
    test_pipeline = cfg.test_dataloader.dataset.pipeline
    data_preprocessor = cfg.model.data_preprocessor
    codec = cfg.codec
    if isinstance(codec, list):
        codec = codec[-1]
    input_size = codec["input_size"] if input_shape is None else input_shape
    test_pipeline[0] = dict(type="LoadImageFromFile")
    for i in reversed(range(len(test_pipeline))):
        trans = test_pipeline[i]
        if trans["type"] == "PackPoseInputs":
            test_pipeline.pop(i)
        elif trans["type"] == "GetBBoxCenterScale":
            trans["type"] = "TopDownGetBboxCenterScale"
            trans["padding"] = 1.25  # default argument
            trans["image_size"] = input_size
        elif trans["type"] == "TopdownAffine":
            trans["type"] = "TopDownAffine"
            trans["image_size"] = input_size
            trans.pop("input_size")

    test_pipeline.append(
        dict(
            type="Normalize",
            mean=data_preprocessor.mean,
            std=data_preprocessor.std,
            to_rgb=data_preprocessor.bgr_to_rgb,
        )
    )
    test_pipeline.append(dict(type="ImageToTensor", keys=["img"]))
    test_pipeline.append(
        dict(
            type="Collect",
            keys=["img"],
            meta_keys=[
                "img_shape",
                "pad_shape",
                "ori_shape",
                "img_norm_cfg",
                "scale_factor",
                "bbox_score",
                "center",
                "scale",
            ],
        )
    )

    cfg.test_dataloader.dataset.pipeline = test_pipeline
    return cfg


def _get_dataset_metainfo(model_cfg: mmengine.Config):
    """Get metainfo of dataset.

    Args:
        model_cfg Config: Input model Config object.
    Returns:
        (list[str], list[np.ndarray]): Class names and palette
    """
    from mmpose import datasets  # noqa
    from mmpose.registry import DATASETS

    module_dict = DATASETS.module_dict

    for dataloader_name in ["test_dataloader", "val_dataloader", "train_dataloader"]:
        if dataloader_name not in model_cfg:
            continue
        dataloader_cfg = model_cfg[dataloader_name]
        dataset_cfg = dataloader_cfg.dataset
        dataset_mmpose = module_dict.get(dataset_cfg.type, None)
        if dataset_mmpose is None:
            continue
        if hasattr(dataset_mmpose, "_load_metainfo") and isinstance(
            dataset_mmpose._load_metainfo, Callable
        ):
            meta = dataset_mmpose._load_metainfo(dataset_cfg.get("metainfo", None))
            if meta is not None:
                return meta
        if hasattr(dataset_mmpose, "METAINFO"):
            return dataset_mmpose.METAINFO

    return None


MMPOSE_TASK = Registry("mmpose_tasks")


@CODEBASE.register_module(Codebase.MMPOSE.value)
class MMPose(MMCodebase):
    """mmpose codebase class."""

    task_registry = MMPOSE_TASK

    @classmethod
    def register_deploy_modules(cls):
        """register rewritings."""
        import juxtapose.mmdeploy.codebase.mmdet.models
        import juxtapose.mmdeploy.codebase.mmdet.ops
        import juxtapose.mmdeploy.codebase.mmdet.structures
        import juxtapose.mmdeploy.codebase.mmpose.models  # noqa: F401

    @classmethod
    def register_all_modules(cls):
        """register all modules from mmpose."""
        from mmpose.utils.setup_env import register_all_modules

        cls.register_deploy_modules()
        register_all_modules(True)


@MMPOSE_TASK.register_module(Task.POSE_DETECTION.value)
class PoseDetection(BaseTask):
    """Pose detection task class.

    Args:
        model_cfg (mmengine.Config): Original PyTorch model config file.
        deploy_cfg (mmengine.Config): Deployment config file or loaded Config
            object.
        device (str): A string represents device type.
    """

    def __init__(
        self, model_cfg: mmengine.Config, deploy_cfg: mmengine.Config, device: str
    ):
        super().__init__(model_cfg, deploy_cfg, device)
        self.model_cfg.model.test_cfg["flip_test"] = False

        from mmcv.transforms import Compose
        from mmpose.registry import TRANSFORMS

        cfg = self.model_cfg

        # build the data pipeline
        test_pipeline = [
            TRANSFORMS.build(c) for c in cfg.test_dataloader.dataset.pipeline
        ]
        self.test_pipeline = Compose(test_pipeline)

    def build_backend_model(
        self,
        model_files: Sequence[str] = None,
        data_preprocessor_updater: Optional[Callable] = None,
        **kwargs,
    ) -> torch.nn.Module:
        """build backend model.

        Args:
            model_files (Sequence[str]): Input model files. Default is None.
            data_preprocessor_updater (Callable | None): A function to update
                the data_preprocessor. Defaults to None.

        Returns:
            nn.Module: An initialized backend model.
        """
        from .pose_detection_model import build_pose_detection_model

        data_preprocessor = self.model_cfg.model.data_preprocessor
        if data_preprocessor_updater is not None:
            data_preprocessor = data_preprocessor_updater(data_preprocessor)
        model = build_pose_detection_model(
            model_files,
            self.model_cfg,
            self.deploy_cfg,
            device=self.device,
            data_preprocessor=data_preprocessor,
            **kwargs,
        )
        return model.eval().to(self.device)

    def create_input(
        self,
        imgs: Union[str, np.ndarray, Sequence],
        input_shape: Sequence[int] = None,
        data_preprocessor: Optional[BaseDataPreprocessor] = None,
        bboxes: Sequence = None,
        **kwargs,
    ) -> Tuple[Dict, torch.Tensor]:
        """Create input for pose detection.

        Args:
            imgs (Any): Input image(s), accepted data type are ``str``,
                ``np.ndarray``.
            input_shape (list[int]): A list of two integer in (width, height)
                format specifying input shape. Defaults to ``None``.
            data_preprocessor (BaseDataPreprocessor | None): Input data pre-
                processor. Default is ``None``.

        Returns:
            tuple: (data, inputs), meta information for the input image
             and input.
        """
        from mmcv.transforms import Compose
        from mmpose.registry import TRANSFORMS

        cfg = self.model_cfg
        if isinstance(imgs, (list, tuple)):
            if not isinstance(imgs[0], (np.ndarray, str)):
                raise AssertionError("imgs must be strings or numpy arrays")
        elif isinstance(imgs, (np.ndarray, str)):
            imgs = [imgs]
            img_path = [imgs]
        else:
            raise AssertionError("imgs must be strings or numpy arrays")
        if isinstance(imgs, (list, tuple)) and isinstance(imgs[0], str):
            img_path = imgs
            img_data = [mmcv.imread(img) for img in imgs]
            imgs = img_data

        if not bboxes:
            person_results = []
            bboxes = []
            for img in imgs:
                height, width = img.shape[:2]
                # create dummy person results
                person_results.append([{"bbox": np.array([0, 0, width, height])}])
                bboxes.append(np.array([box["bbox"] for box in person_results[-1]]))

        if input_shape is not None and hasattr(cfg, "codec"):
            if isinstance(cfg.codec, dict):
                codec = cfg.codec
            elif isinstance(cfg.codec, list):
                codec = cfg.codec[0]
            else:
                raise TypeError(f"Unsupported type {type(cfg.codec)}")
            input_size = codec["input_size"]
            if tuple(input_shape) != tuple(input_size):
                logger = get_root_logger()
                logger.warning(
                    f"Input shape from deploy config is not "
                    f"same as input_size in model config:"
                    f"{input_shape} vs {input_size}"
                )
        batch_data = defaultdict(list)
        meta_data = _get_dataset_metainfo(self.model_cfg)

        if bboxes:
            for i in range(len(bboxes)):
                for bbox in bboxes[i]:
                    # prepare data
                    bbox_score = np.array(
                        [bbox[4] if len(bbox) == 5 else 1]
                    )  # shape (1,)
                    data = {
                        "img": imgs[0],
                        "bbox_score": bbox_score,
                        "bbox": []
                        if hasattr(cfg.model, "bbox_head")
                        and cfg.model.bbox_head.type == "YOLOXPoseHead"
                        else bbox[None],
                        "img_path": img_path[0],
                    }
                    data.update(meta_data)
                    data = self.test_pipeline(data)
                    data["inputs"] = data["inputs"].to(self.device)
                    batch_data["inputs"].append(data["inputs"])
                    batch_data["data_samples"].append(data["data_samples"])

        else:
            assert len(imgs) == len(bboxes) == len(person_results)
            for i in range(len(imgs)):
                for bbox in bboxes[i]:
                    # prepare data
                    bbox_score = np.array(
                        [bbox[4] if len(bbox) == 5 else 1]
                    )  # shape (1,)
                    data = {
                        "img": imgs[i],
                        "bbox_score": bbox_score,
                        "bbox": []
                        if hasattr(cfg.model, "bbox_head")
                        and cfg.model.bbox_head.type == "YOLOXPoseHead"
                        else bbox[None],
                        "img_path": img_path[i],
                    }
                    data.update(meta_data)
                    data = self.test_pipeline(data)
                    data["inputs"] = data["inputs"].to(self.device)
                    batch_data["inputs"].append(data["inputs"])
                    batch_data["data_samples"].append(data["data_samples"])

        if data_preprocessor is not None:
            batch_data = data_preprocessor(batch_data, False)
        input_tensor = batch_data["inputs"]
        return batch_data, input_tensor

    def visualize(
        self,
        image: Union[str, np.ndarray],
        result: list,
        output_file: str,
        window_name: str,
        show_result: bool = False,
        **kwargs,
    ):
        """Visualize predictions of a model.

        Args:
            image (str | np.ndarray): Input image to draw predictions on.
            result (list): A list of predictions.
            output_file (str): Output file to save drawn image.
            window_name (str): The name of visualization window. Defaults to
                an empty string.
            show_result (bool): Whether to show result in windows, defaults
                to `False`.
        """
        from mmpose.apis.inference import dataset_meta_from_config
        from mmpose.visualization import PoseLocalVisualizer

        save_dir, filename = os.path.split(output_file)
        name = os.path.splitext(filename)[0]
        dataset_meta = dataset_meta_from_config(self.model_cfg, dataset_mode="test")
        visualizer = PoseLocalVisualizer(name=name, save_dir=save_dir)
        visualizer.set_dataset_meta(dataset_meta)

        if isinstance(image, str):
            image = mmcv.imread(image, channel_order="rgb")
        draw_bbox = result.pred_instances.bboxes is not None
        if draw_bbox and isinstance(result.pred_instances.bboxes, torch.Tensor):
            result.pred_instances.bboxes = result.pred_instances.bboxes.cpu().numpy()
        visualizer.add_datasample(
            name,
            image,
            data_sample=result,
            draw_gt=False,
            draw_bbox=draw_bbox,
            show=show_result,
            out_file=output_file,
        )

    def get_model_name(self, *args, **kwargs) -> str:
        """Get the model name.

        Return:
            str: the name of the model.
        """
        assert "type" in self.model_cfg.model, "model config contains no type"
        name = self.model_cfg.model.type.lower()
        return name

    @staticmethod
    def get_partition_cfg(partition_type: str, **kwargs) -> Dict:
        """Get a certain partition config for mmpose.

        Args:
            partition_type (str): A string specifying partition type.
        """
        raise NotImplementedError("Not supported yet.")

    def get_preprocess(self, *args, **kwargs) -> Dict:
        """Get the preprocess information for SDK.

        Return:
            dict: Composed of the preprocess information.
        """
        input_shape = get_input_shape(self.deploy_cfg)
        model_cfg = process_model_config(self.model_cfg, [""], input_shape)
        preprocess = model_cfg.test_dataloader.dataset.pipeline
        return preprocess

    def get_postprocess(self, *args, **kwargs) -> Dict:
        """Get the postprocess information for SDK."""
        codec = self.model_cfg.codec
        if isinstance(codec, (list, tuple)):
            codec = codec[-1]
        component = "UNKNOWN"
        params = copy.deepcopy(self.model_cfg.model.test_cfg)
        params.update(codec)
        if self.model_cfg.model.type == "TopdownPoseEstimator":
            component = "TopdownHeatmapSimpleHeadDecode"
            if codec.type == "MSRAHeatmap":
                params["post_process"] = "default"
            elif codec.type == "UDPHeatmap":
                params["post_process"] = "default"
                params["use_udp"] = True
            elif codec.type == "MegviiHeatmap":
                params["post_process"] = "megvii"
                params["modulate_kernel"] = self.model_cfg.kernel_sizes[-1]
            elif codec.type == "SimCCLabel":
                component = "SimCCLabelDecode"
            elif codec.type == "RegressionLabel":
                component = "DeepposeRegressionHeadDecode"
            elif codec.type == "IntegralRegressionLabel":
                component = "DeepposeRegressionHeadDecode"
            else:
                raise RuntimeError(f"Unsupported codecs type: {codec.type}")
        postprocess = dict(params=params, type=component)
        return postprocess
