# 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.

import copy
import logging
from typing import Dict, List, Optional, Sequence, Union

from datumaro import Image
from datumaro.components.annotation import Bbox, Polygon

from geti_sdk.annotation_readers.base_annotation_reader import AnnotationReader
from geti_sdk.data_models import Annotation as SCAnnotation
from geti_sdk.data_models import TaskType
from geti_sdk.data_models.enums.task_type import GLOBAL_TASK_TYPES
from geti_sdk.data_models.media import MediaInformation
from geti_sdk.rest_converters import AnnotationRESTConverter
from geti_sdk.utils import generate_segmentation_labels, get_dict_key_from_value

from .datumaro_dataset import DatumaroDataset


class DatumAnnotationReader(AnnotationReader):
    """
    Class to read annotations using datumaro
    """

    _SUPPORTED_TASK_TYPES = [
        TaskType.DETECTION,
        TaskType.SEGMENTATION,
        TaskType.CLASSIFICATION,
        TaskType.INSTANCE_SEGMENTATION,
        TaskType.ROTATED_DETECTION,
        TaskType.ANOMALY_CLASSIFICATION,
        TaskType.ANOMALY_DETECTION,
        TaskType.ANOMALY_SEGMENTATION,
    ]

    def __init__(
        self,
        base_data_folder: str,
        annotation_format: str,
        task_type: Union[TaskType, str] = TaskType.DETECTION,
    ):
        super().__init__(
            base_data_folder=base_data_folder,
            annotation_format=annotation_format,
            task_type=task_type,
        )
        self.dataset = DatumaroDataset(
            dataset_format=annotation_format, dataset_path=base_data_folder
        )
        self._override_label_map: Optional[Dict[int, str]] = None
        self._applied_filters: List[Dict[str, Union[List[str], str]]] = []

    def prepare_and_set_dataset(
        self,
        task_type: Union[TaskType, str],
        previous_task_type: Optional[TaskType] = None,
    ) -> None:
        """
        Prepare the dataset for a specific `task_type`. This could involve for
        example conversion of annotation shapes.

        :param task_type: TaskType for which to prepare the dataset.
        :param previous_task_type: TaskType preceding the task to prepare the dataset
            for
        """
        if not isinstance(task_type, TaskType):
            task_type = TaskType(task_type)
        if task_type != self.task_type:
            logging.info(f"Task type changed to {task_type} for dataset")
            if task_type not in self._SUPPORTED_TASK_TYPES:
                raise ValueError(f"Unsupported task type {task_type}")
            new_dataset = DatumaroDataset(
                dataset_format=self.annotation_format, dataset_path=self.base_folder
            )
            self.task_type = task_type
            self.dataset = new_dataset
            for filter_parameters in self.applied_filters:
                self.filter_dataset(**filter_parameters)

        dataset = self.dataset.prepare_dataset(
            task_type=task_type, previous_task_type=previous_task_type
        )
        self.dataset.set_dataset(dataset)
        logging.info(f"Dataset is prepared for {task_type} task.")

    def convert_labels_to_segmentation_names(self) -> None:
        """
        Convert the label names in a dataset to '*_shape`, where `*` is
        the original label name. It can be used to generate unique label names for the
        segmentation task in a detection_to_segmentation project
        """
        segmentation_label_map: Dict[int, str] = {}
        label_names = list(self.datum_label_map.values())
        segmentation_label_names = generate_segmentation_labels(label_names)
        for datum_index, label_name in self.datum_label_map.items():
            label_index = label_names.index(label_name)
            segmentation_label_map.update(
                {datum_index: segmentation_label_names[label_index]}
            )
        self.override_label_map(segmentation_label_map)

    def get_all_label_names(self) -> List[str]:
        """
        Retrieve the list of labels names from a datumaro dataset.
        """
        return list(set(self.datum_label_map.values()))

    @property
    def datum_label_map(self) -> Dict[int, str]:
        """
        :return: Dictionary mapping the datumaro label id to the label name
        """
        if self._override_label_map is None:
            return self.dataset.label_mapping
        else:
            return self._override_label_map

    def override_label_map(self, new_label_map: Dict[int, str]):
        """
        Override the label map defined in the datumaro dataset
        """
        self._override_label_map = new_label_map

    def reset_label_map(self):
        """
        Reset the label map back to the original one from the datumaro dataset.
        """
        self._override_label_map = None

    def get_all_image_names(self) -> List[str]:
        """
        Return a list of image names in the dataset
        """
        return self.dataset.image_names

    def get_data(
        self,
        filename: str,
        label_name_to_id_mapping: dict,
        media_information: MediaInformation,
        preserve_shape_for_global_labels: bool = False,
    ) -> List[SCAnnotation]:
        """
        Return the annotation data for the dataset item corresponding to `filename`.

        :param filename: name of the item to get the annotation data for.
        :param label_name_to_id_mapping: mapping of label name to label id.
        :param media_information: MediaInformation object containing information
            (e.g. width, height) about the media item to upload the annotation for
        :param preserve_shape_for_global_labels: False to convert shapes for global
            tasks to full rectangles (required for classification like tasks in
            Intel® Geti™ projects), True to preserve such shapes. This parameter
            should be:

             - False when uploading annotations to a single task project
             - True when uploading annotations for a classification like task,
                following a local task in a task chain project.

        :return: List of Annotation objects containing all annotations for the given
            dataset item.
        """
        ds_item = self.dataset.get_item_by_id(filename)
        image_size = ds_item.media_as(Image).size
        annotation_list: List[SCAnnotation] = []
        labels = []

        # Remove duplicate annotations, datumaro does not check for this
        datum_annotations = [
            i
            for n, i in enumerate(ds_item.annotations)
            if i not in ds_item.annotations[:n]
        ]

        for annotation in datum_annotations:
            try:
                label_name = self.datum_label_map[annotation.label]
            except KeyError:
                # Label is not in the Intel® Geti™ project labels, move on to next
                # annotation for this dataset item.
                continue

            label_id = label_name_to_id_mapping.get(label_name)
            label = {"id": label_id, "probability": 1.0}
            if (
                self.task_type not in GLOBAL_TASK_TYPES
                or preserve_shape_for_global_labels
            ):
                if isinstance(annotation, Bbox):
                    x1 = float(annotation.points[0])
                    y1 = float(annotation.points[1])
                    x2 = float(annotation.points[2])
                    y2 = float(annotation.points[3])
                    shape = {
                        "type": "RECTANGLE",
                        "x": x1,
                        "y": y1,
                        "width": x2 - x1,
                        "height": y2 - y1,
                    }
                elif isinstance(annotation, Polygon):
                    points = [
                        {"x": float(x), "y": float(y)}
                        for x, y in zip(*[iter(annotation.points)] * 2)
                    ]
                    shape = {"type": "POLYGON", "points": points}
                else:
                    logging.warning(
                        f"Unsupported annotation type found: "
                        f"{type(annotation)}. Skipping..."
                    )
                    continue
                sc_annotation = AnnotationRESTConverter.annotation_from_dict(
                    {"labels": [label], "shape": shape}
                )
                annotation_list.append(sc_annotation)
            else:
                labels.append(label)

        if not preserve_shape_for_global_labels and self.task_type in GLOBAL_TASK_TYPES:
            shape = {
                "type": "RECTANGLE",
                "x": 0.0,
                "y": 0.0,
                "width": float(image_size[1]),
                "height": float(image_size[0]),
            }
            sc_annotation = AnnotationRESTConverter.annotation_from_dict(
                {"labels": labels, "shape": shape}
            )
            annotation_list.append(sc_annotation)
        return annotation_list

    @property
    def applied_filters(self) -> List[Dict[str, Union[List[str], str]]]:
        """
        Return a list of filters and their parameters that have been previously
        applied to the dataset.
        """
        return copy.deepcopy(self._applied_filters)

    def filter_dataset(self, labels: Sequence[str], criterion="OR") -> None:
        """
        Retain only those items with annotations in the list of labels passed.

        :param: labels     List of labels to filter on
        :param: criterion  Filter criterion, currently "OR" or "AND" are implemented
        """
        self.dataset.filter_items_by_labels(labels=labels, criterion=criterion)
        self._applied_filters.append({"labels": labels, "criterion": criterion})

    def group_labels(self, labels_to_group: List[str], group_name: str) -> None:
        """
        Group multiple labels into one. Grouping converts the list of labels into one
        single label named `group_name`.

        This method does not return anything, but instead overrides the label map for
        the datamaro dataset to account for the grouping.

        :param labels_to_group: List of labels names that should be grouped together
        :param group_name: Name of the resulting label
        :return:
        """
        label_keys = [
            get_dict_key_from_value(self.datum_label_map, label)
            for label in labels_to_group
        ]
        new_label_map = copy.deepcopy(self.datum_label_map)
        for key in label_keys:
            new_label_map[key] = group_name
        self.override_label_map(new_label_map=new_label_map)

    def get_annotation_stats(self) -> Dict[str, Dict[str, int]]:
        """
        Return the object and image counts per label in the dataset.

        :return: Dictionary containing label names as keys, and as values:
            - n_images: Number of images containing this label
            - n_objects: Number of independent objects with this label
        """
        label_statistics: Dict[str, Dict[str, int]] = {}
        for item in self.dataset.dataset:
            label_names = set(
                [
                    self.datum_label_map[annotation.label]
                    for annotation in item.annotations
                ]
            )
            object_counts: Dict[str, int] = {}
            for label in label_names:
                object_counts[label] = 0
            for annotation in item.annotations:
                label_name = self.datum_label_map[annotation.label]
                object_counts[label_name] += 1

            for label in label_names:
                label_stats = label_statistics.get(label, {})
                if not label_stats:
                    label_statistics[label] = {
                        "n_images": 1,
                        "n_objects": object_counts[label],
                    }
                else:
                    label_statistics[label]["n_images"] += 1
                    label_statistics[label]["n_objects"] += object_counts[label]
        return label_statistics
