# 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 glob
import json
import logging
import os
import warnings
from random import sample
from typing import Any, Dict, List, Optional, Union

from geti_sdk.data_models import Annotation, TaskType
from geti_sdk.data_models.media import MediaInformation
from geti_sdk.rest_converters import AnnotationRESTConverter
from geti_sdk.rest_converters.annotation_rest_converter import (
    NormalizedAnnotationRESTConverter,
)

from .base_annotation_reader import AnnotationReader


class GetiAnnotationReader(AnnotationReader):
    """
    AnnotationReader for loading annotation files in Intel® Geti™ format.
    """

    def __init__(
        self,
        base_data_folder: str,
        annotation_format: str = ".json",
        task_type: Optional[Union[TaskType, str]] = None,
        label_names_to_include: Optional[List[str]] = None,
    ):
        """
        :param base_data_folder: Path to the folder containing the annotations
        :param annotation_format: Extension of the annotation files. Defaults to ".json"
        :param task_type: Optional task type to prepare the annotations for. Can also
            be specified later.
        :param label_names_to_include: Names of the labels that should be included
            when reading annotation data. This can be used to filter the annotations
            for certain labels.
        """
        if annotation_format != ".json":
            raise ValueError(
                f"Annotation format {annotation_format} is currently not"
                f" supported by the GetiAnnotationReader"
            )
        super().__init__(
            base_data_folder=base_data_folder,
            annotation_format=annotation_format,
            task_type=task_type,
        )
        self._label_names_to_include = label_names_to_include
        self._normalized_annotations = self._has_normalized_annotations()

    def _get_label_names(self, all_labels: List[str]) -> List[str]:
        """
        Return the labels for the task type the annotation reader is currently set to.

        :param all_labels: List of all label names in the project
        """
        if self._label_names_to_include is None:
            logging.info("No label mapping defined, including all labels")
            labels = all_labels
        else:
            labels = [
                label for label in all_labels if label in self._label_names_to_include
            ]
        return labels

    def _get_raw_annotation_data(self, filename: str) -> Dict[str, Any]:
        """
        Read the annotation data from the file at `filename`

        :param filename: Name of the annotation file to read
        :return: Dictionary holding the annotation data
        """
        filepath = glob.glob(
            os.path.join(self.base_folder, f"{filename}{self.annotation_format}"),
        )
        if len(filepath) > 1:
            warnings.warn(
                f"Multiple matching annotation files found for image with "
                f"name {filename}. Skipping this image..."
            )
            data = {"annotations": []}
        elif len(filepath) == 0:
            logging.info(
                f"No matching annotation file found for image with name {filename}."
                f" Skipping this image..."
            )
            data = {"annotations": []}
        else:
            filepath = filepath[0]
            with open(filepath, "r") as f:
                data = json.load(f)
        return data

    def get_data(
        self,
        filename: str,
        label_name_to_id_mapping: dict,
        media_information: MediaInformation,
        preserve_shape_for_global_labels: bool = False,
    ) -> List[Annotation]:
        """
        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.
        """
        data = self._get_raw_annotation_data(filename=filename)

        new_annotations = []
        for annotation in data["annotations"]:
            if self._normalized_annotations:
                annotation_object = (
                    NormalizedAnnotationRESTConverter.normalized_annotation_from_dict(
                        annotation,
                        image_width=media_information.width,
                        image_height=media_information.height,
                    )
                )
            else:
                annotation_object = AnnotationRESTConverter.annotation_from_dict(
                    annotation
                )
            for label in annotation_object.labels:
                label.id = label_name_to_id_mapping[label.name]
            for label_dict in annotation["labels"]:
                if self.task_type is not None:
                    if label_dict["name"] not in self._get_label_names(
                        list(label_name_to_id_mapping.keys())
                    ):
                        annotation_object.pop_label_by_name(
                            label_name=label_dict["name"]
                        )
            new_annotations.append(annotation_object)
        return new_annotations

    def get_all_label_names(self) -> List[str]:
        """
        Retrieve the unique label names for all annotations in the annotation folder

        :return: List of label names
        """
        logging.info(f"Reading annotation files in folder {self.base_folder}...")
        unique_label_names = []
        annotation_files = glob.glob(
            os.path.join(self.base_folder, f"*{self.annotation_format}"),
        )
        if len(annotation_files) == 0:
            raise ValueError(
                f"No valid annotation files were found in folder {self.base_folder}"
            )
        for annotation_file in annotation_files:
            with open(annotation_file, "r") as f:
                data = json.load(f)
            annotations = data.get("annotations", None)
            if annotations is None:
                raise ValueError(
                    f"Annotation file '{annotation_file}' does not contain any "
                    f"annotations. Please make sure that this is a valid "
                    f"annotation file."
                )
            for annotation in annotations:
                labels = [label["name"] for label in annotation["labels"]]
                for label in labels:
                    if label not in unique_label_names:
                        unique_label_names.append(label)
        return unique_label_names

    def _has_normalized_annotations(self) -> bool:
        """
        Check if the annotation files belonging to this annotation reader are normalized
        """
        filenames = self.get_data_filenames()
        n_sample = min(len(filenames), 50)
        if n_sample == 50:
            list_to_check = sample(filenames, n_sample)
        else:
            list_to_check = filenames

        NORMALIZED_KEY = "normalized"
        PIXEL_KEY = "pixel"

        annotation_stats = {NORMALIZED_KEY: 0, PIXEL_KEY: 0}
        for filename in list_to_check:
            data = self._get_raw_annotation_data(filename=filename)
            for annotation_dict in data["annotations"]:
                annotation_object = AnnotationRESTConverter.annotation_from_dict(
                    annotation_dict
                )
                shape = annotation_object.shape
                x_max, y_max = shape.x_max, shape.y_max
                if x_max <= 1 and y_max <= 1:
                    annotation_stats[NORMALIZED_KEY] += 1
                else:
                    annotation_stats[PIXEL_KEY] += 1

        if annotation_stats[NORMALIZED_KEY] == 0:
            return False
        elif annotation_stats[PIXEL_KEY] == 0:
            logging.info(
                "Legacy annotation format detected. The annotations you are trying to "
                "upload were most likely downloaded from a pre-production version of "
                "the Intel Geti software. They will be converted to the latest "
                "annotation format upon upload to the Intel Geti platform. "
            )
            return True
        else:
            raise ValueError(
                f"The annotation directory '{self.base_folder}' contains both "
                f"normalized ({annotation_stats[NORMALIZED_KEY]} shapes) and "
                f"non-normalized ({annotation_stats[PIXEL_KEY]} shapes) objects. "
                f"Unable to parse annotation data."
            )
