"""OTXMaskRCNNModel & OTXSSDModel of OTX Detection."""

# Copyright (C) 2022 Intel Corporation
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions
# and limitations under the License.

from typing import Dict

import numpy as np

try:
    from openvino.model_zoo.model_api.models.instance_segmentation import MaskRCNNModel
    from openvino.model_zoo.model_api.models.ssd import SSD, find_layer_by_name
    from openvino.model_zoo.model_api.models.utils import Detection
except ImportError as e:
    import warnings

    warnings.warn(f"{e}: ModelAPI was not found.")


class OTXMaskRCNNModel(MaskRCNNModel):
    """OpenVINO model wrapper for OTX MaskRCNN model."""

    __model__ = "OTX_MaskRCNN"

    def _get_outputs(self):
        output_match_dict = {}
        output_names = ["boxes", "labels", "masks", "feature_vector", "saliency_map"]
        for output_name in output_names:
            for node_name, node_meta in self.outputs.items():
                if output_name in node_meta.names:
                    output_match_dict[output_name] = node_name
                    break
        return output_match_dict

    def postprocess(self, outputs, meta):
        """Post process function for OTX MaskRCNN model."""

        # pylint: disable-msg=too-many-locals
        resize_mask = meta.get("resize_mask", True)

        # FIXME: here, batch dim of IR must be 1
        boxes = outputs[self.output_blob_name["boxes"]]
        boxes = boxes.squeeze(0)
        assert boxes.ndim == 2
        masks = outputs[self.output_blob_name["masks"]]
        masks = masks.squeeze(0)
        assert masks.ndim == 3
        classes = outputs[self.output_blob_name["labels"]].astype(np.uint32)
        classes = classes.squeeze(0)
        assert classes.ndim == 1
        if self.is_segmentoly:
            scores = outputs[self.output_blob_name["scores"]]
        else:
            scores = boxes[:, 4]
            boxes = boxes[:, :4]
            classes += 1

        # Filter out detections with low confidence.
        detections_filter = scores > self.confidence_threshold  # pylint: disable=no-member
        scores = scores[detections_filter]
        boxes = boxes[detections_filter]
        masks = masks[detections_filter]
        classes = classes[detections_filter]

        scale_x = meta["resized_shape"][1] / meta["original_shape"][1]
        scale_y = meta["resized_shape"][0] / meta["original_shape"][0]
        boxes[:, 0::2] /= scale_x
        boxes[:, 1::2] /= scale_y

        resized_masks = []
        for box, cls, raw_mask in zip(boxes, classes, masks):
            raw_cls_mask = raw_mask[cls, ...] if self.is_segmentoly else raw_mask
            if resize_mask:
                resized_masks.append(self._segm_postprocess(box, raw_cls_mask, *meta["original_shape"][:-1]))
            else:
                resized_masks.append(raw_cls_mask)

        return scores, classes, boxes, resized_masks

    def segm_postprocess(self, *args, **kwargs):
        """Post-process for segmentation masks."""
        return self._segm_postprocess(*args, **kwargs)


class OTXSSDModel(SSD):
    """OpenVINO model wrapper for OTX SSD model."""

    __model__ = "OTX_SSD"

    def __init__(self, model_adapter, configuration=None, preload=False):
        # pylint: disable-next=bad-super-call
        super(SSD, self).__init__(model_adapter, configuration, preload)
        self.image_info_blob_name = self.image_info_blob_names[0] if len(self.image_info_blob_names) == 1 else None
        self.output_parser = BatchBoxesLabelsParser(
            self.outputs,
            self.inputs[self.image_blob_name].shape[2:][::-1],
        )

    def _get_outputs(self) -> Dict:
        """Match the output names with graph node index."""
        output_match_dict = {}
        output_names = ["boxes", "labels", "feature_vector", "saliency_map"]
        for output_name in output_names:
            for node_name, node_meta in self.outputs.items():
                if output_name in node_meta.names:
                    output_match_dict[output_name] = node_name
                    break
        return output_match_dict


class BatchBoxesLabelsParser:
    """Batched output parser."""

    def __init__(self, layers, input_size, labels_layer="labels", default_label=0):
        try:
            self.labels_layer = find_layer_by_name(labels_layer, layers)
        except ValueError:
            self.labels_layer = None
            self.default_label = default_label

        self.bboxes_layer = self.find_layer_bboxes_output(layers)
        self.input_size = input_size

    @staticmethod
    def find_layer_bboxes_output(layers):
        """find_layer_bboxes_output."""
        filter_outputs = [name for name, data in layers.items() if len(data.shape) == 3 and data.shape[-1] == 5]
        if not filter_outputs:
            raise ValueError("Suitable output with bounding boxes is not found")
        if len(filter_outputs) > 1:
            raise ValueError("More than 1 candidate for output with bounding boxes.")
        return filter_outputs[0]

    def __call__(self, outputs):
        """Parse bboxes."""
        # FIXME: here, batch dim of IR must be 1
        bboxes = outputs[self.bboxes_layer]
        bboxes = bboxes.squeeze(0)
        assert bboxes.ndim == 2
        scores = bboxes[:, 4]
        bboxes = bboxes[:, :4]
        bboxes[:, 0::2] /= self.input_size[0]
        bboxes[:, 1::2] /= self.input_size[1]
        if self.labels_layer:
            labels = outputs[self.labels_layer]
        else:
            labels = np.full(len(bboxes), self.default_label, dtype=bboxes.dtype)
        labels = labels.squeeze(0)

        detections = [Detection(*bbox, score, label) for label, score, bbox in zip(labels, scores, bboxes)]
        return detections
