# Copyright (c) Alibaba, Inc. and its affiliates.

import os
from typing import Any, Dict

import numpy as np
import torch

from modelscope.metainfo import Pipelines
from modelscope.outputs import OutputKeys
from modelscope.pipelines.base import Input, Pipeline
from modelscope.pipelines.builder import PIPELINES
from modelscope.preprocessors import LoadImage
from modelscope.utils.constant import ModelFile, Tasks


@PIPELINES.register_module(
    Tasks.image_fewshot_detection,
    module_name=Pipelines.image_fewshot_detection)
class ImageDefrcnDetectionPipeline(Pipeline):
    """ Image DeFRCN few-shot detection Pipeline. Given a image,
        pipeline will return the detection results on the image.
        Example:

        ```python
        >>> from modelscope.pipelines import pipeline
        >>> detector = pipeline('image-fewshot-detection', 'damo/cv_resnet101_detection_fewshot-defrcn')
        >>> detector('/Path/Image')
           {
            'scores': [0.8307567834854126, 0.1606406420469284],
            'labels': ['person', 'dog'],
            'boxes': [
                [27.391937255859375, 0.0, 353.0, 500.0],
                [64.22428131103516, 229.2884521484375, 213.90573120117188, 370.0657958984375]
            ]
            }
        >>> #
        ```
    """

    def __init__(self, model: str, **kwargs):
        """
            model: model id on modelscope hub.
        """
        super().__init__(model=model, auto_collate=False, **kwargs)

        model_path = os.path.join(self.model.model_dir,
                                  ModelFile.TORCH_MODEL_FILE)
        self.model.model = self._load_pretrained(
            self.model.model, model_path, self.model.model_cfg.MODEL.DEVICE)

    def _load_pretrained(self, net, load_path, device='cuda', strict=True):

        load_net = torch.load(load_path, map_location=device)
        if 'scheduler' in load_net:
            del load_net['scheduler']
        if 'optimizer' in load_net:
            del load_net['optimizer']
        if 'iteration' in load_net:
            del load_net['iteration']
        net.load_state_dict(load_net['model'], strict=strict)

        return net

    def preprocess(self, input: Input) -> Dict[str, Any]:

        img = LoadImage.convert_to_ndarray(input)
        img = img.astype(np.float)

        image = img[..., ::-1].copy()  # rgb to bgr
        tim = torch.Tensor(image).permute(2, 0, 1)

        result = {'image': tim}
        return result

    def forward(self, input: Dict[str, Any]) -> Dict[str, Any]:

        outputs = self.model.inference(input)
        result = {'data': outputs}
        return result

    def postprocess(self, inputs: Dict[str, Any]) -> Dict[str, Any]:
        if inputs['data'] is None:
            outputs = {
                OutputKeys.SCORES: [],
                OutputKeys.LABELS: [],
                OutputKeys.BOXES: []
            }
            return outputs

        objects = inputs['data']['instances'].get_fields()
        labels, bboxes = [], []
        for label, box in zip(objects['pred_classes'], objects['pred_boxes']):
            labels.append(self.model.config.model.classes[label])
            bboxes.append(box.tolist())

        scores = objects['scores'].tolist()

        outputs = {
            OutputKeys.SCORES: scores,
            OutputKeys.LABELS: labels,
            OutputKeys.BOXES: bboxes
        }
        return outputs
