# Copyright Open Logistics Foundation
#
# Licensed under the Open Logistics Foundation License 1.3.
# For details on the licensing terms, see the LICENSE file.
# SPDX-License-Identifier: OLFL-1.3

"""Module for parsing Label Studio formatted annotations of single json files that are
generated by Label Studio per image
"""
import copy
import json
import logging
import os
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple, cast

import cv2
import numpy as np
from attr import define

from mlcvzoo_base.api.data.annotation import BaseAnnotation
from mlcvzoo_base.api.data.annotation_class_mapper import AnnotationClassMapper
from mlcvzoo_base.api.data.annotation_parser import AnnotationParser
from mlcvzoo_base.api.data.bounding_box import BoundingBox
from mlcvzoo_base.api.data.box import Box, rotate_point
from mlcvzoo_base.api.data.class_identifier import ClassIdentifier
from mlcvzoo_base.api.data.classification import Classification
from mlcvzoo_base.api.data.segmentation import Segmentation
from mlcvzoo_base.api.data.types import FrameShape, Point2f
from mlcvzoo_base.api.exceptions import ClassMappingNotFoundError, ForbiddenClassError
from mlcvzoo_base.configuration.annotation_handler_config import (
    AnnotationHandlerLabelStudioInputDataConfig,
)
from mlcvzoo_base.configuration.structs import ObjectDetectionBBoxFormats
from mlcvzoo_base.utils.file_utils import get_image_file_dict

logger = logging.getLogger(__name__)


@define
class LabelStudioAnnotationData:
    image_shape: Optional[Tuple[int, int]] = None
    classifications: List[Classification] = []
    bounding_boxes: List[BoundingBox] = []
    segmentations: List[Segmentation] = []


def __merge_single_entry(results: Dict[Any, Any]) -> Dict[Any, Any]:
    """
    Merge the annotation data of a label studio result entry.
    Typically, label studio creates multiple different entries
    per result type. To parse the annotation content into a
    mlcvzoo annotation we want all the different entries to be
    merged into one result entry that contains all the necessary
    information.

    Args:
        results: The results that should be merged

    Returns:
        The merged results
    """
    merged_results = {}
    for result in results:
        result_id = result["id"]
        result_type = result["type"]
        if result_id not in merged_results:
            merged_results[result_id] = result
            merged_results[result_id]["type"] = [result_type]
            continue

        merged_results[result_id]["value"].update(result["value"])
        merged_results[result_id]["type"] += [result_type]

    return merged_results


def get_label_studio_annotation_dict(
    label_studio_annotation_path: str,
) -> Optional[Dict[Any, Any]]:
    """
    Read the content of a label studio annotation file in json format.
    These files are generated by label studio per annotation task (image).

    If the given label_studio_annotation_path exists, the json data will
    be parsed and merged as preparation for the main parsing methods
    that will translate the data into mlcvzoo annotations.

    Args:
        label_studio_annotation_path: The path to the annotation file

    Returns:
        The dictionary with the parsed annotation data
    """
    with open(label_studio_annotation_path, "r") as label_studio_file:
        try:
            annotation_data = cast(Dict[Any, Any], json.load(label_studio_file))

            annotation_data_orig = copy.deepcopy(annotation_data)

            try:
                # list() is needed for correct type conversion from dict_values to list
                annotation_data["result"] = list(
                    __merge_single_entry(annotation_data["result"]).values()
                )
                return annotation_data
            except KeyError as error:
                # This is not an exception that is breaking the code.
                # We can still use the original parsed data to continue with.
                logger.warning(
                    "Could not merge data due to error=%s, "
                    "will continue with original data" % repr(error)
                )
                return annotation_data_orig

        except (KeyError, json.decoder.JSONDecodeError) as error:
            logger.exception(
                msg="Could not parse annotation file '%s'. Annotation will be skipped."
                % label_studio_annotation_path,
                exc_info=error,
            )

    return None


def get_label_studio_annotation_dict_single_file(
    label_studio_annotation_path: str,
) -> Optional[List[Dict[str, Any]]]:
    """
    Read the content of a label studio annotation file in json format.
    This is expected to be an annotation file that has been exported
    via the label studio UI and contains annotations of multiple annotation tasks.

    If the given label_studio_annotation_path exists, the json data will
    be parsed and merged as preparation for the main parsing methods
    that will translate the data into mlcvzoo annotations.

    Args:
        label_studio_annotation_path: The path to the annotation file

    Returns:
        The dictionary with the parsed annotation data
    """

    with open(label_studio_annotation_path, "r") as label_studio_file:
        try:
            annotation_data = cast(List[Dict[str, Any]], json.load(label_studio_file))

            annotation_data_orig = copy.deepcopy(annotation_data)

            try:
                for index, project in enumerate(annotation_data):
                    annotations = project["annotations"]

                    for annotation_index, annotation in enumerate(annotations):
                        # list() is needed for correct type conversion from dict_values to list
                        annotation_data[index]["annotations"][annotation_index]["result"] = list(
                            __merge_single_entry(
                                annotation_data[index]["annotations"][annotation_index]["result"]
                            ).values()
                        )
            except KeyError as error:
                # This is not an exception that is breaking the code.
                # We can still use the original parsed data to continue with.
                logger.warning(
                    "Could not merge data due to error=%s, "
                    "will continue with original data" % repr(error)
                )
                return annotation_data_orig

            return annotation_data

        except json.decoder.JSONDecodeError as json_error:
            logger.exception(
                msg="Could not parse annotation file '%s'. Annotation will be skipped."
                % label_studio_annotation_path,
                exc_info=json_error,
            )

    return None


def parse_classification(
    mapper: AnnotationClassMapper,
    label_studio_value: Dict[str, Any],
) -> Optional[Classification]:
    """
    Try to parse a single Classification object from a given label studio result value (one
    entry of a result in the annotation file).

    Args:
        mapper: The mapper that should be used to translate class names
        label_studio_value: The input value for parsing the BoundingBox

    Returns:
        Optional[Classification]: The parsed Classification object
    """
    choices = label_studio_value.get("choices", None)
    if choices:
        # Currently we only support single class training, therefore take the first
        # choice occurrence
        label_studio_class_name = choices[0]
    else:
        # We don't accept annotations without a class
        return None

    try:
        class_name = mapper.map_annotation_class_name_to_model_class_name(
            class_name=label_studio_class_name
        )
        class_id = mapper.map_annotation_class_name_to_model_class_id(
            class_name=label_studio_class_name
        )
    except ClassMappingNotFoundError:
        logger.debug(
            "Could not find a valid class-mapping for class-name '%s'. "
            "Classification will be skipped",
            label_studio_class_name,
        )
        return None

    return Classification(
        class_identifier=ClassIdentifier(class_id=class_id, class_name=class_name),
        score=1.0,
    )


def parse_bounding_box(
    mapper: AnnotationClassMapper, label_studio_value: Dict[str, Any], image_shape: Tuple[int, int]
) -> Optional[BoundingBox]:
    """
    Try to parse a single BoundingBox object from a given label studio result value (one
    entry of a result in the annotation file).

    Args:
        mapper: The mapper that should be used to translate class names
        label_studio_value: The input value for parsing the BoundingBox
        image_shape: The shape of the related image

    Returns:
        Optional[BoundingBox]: The parsed BoundingBox object
    """
    rectanglelabels = label_studio_value.get("rectanglelabels", None)
    labels = label_studio_value.get("labels", None)
    if rectanglelabels:
        label_studio_class_name = rectanglelabels[0]
    elif labels:
        label_studio_class_name = labels[0]
    else:
        # We don't except annotations without a class
        return None

    try:
        class_name = mapper.map_annotation_class_name_to_model_class_name(
            class_name=label_studio_class_name
        )
        class_id = mapper.map_annotation_class_name_to_model_class_id(
            class_name=label_studio_class_name
        )
    except ClassMappingNotFoundError:
        logger.debug(
            "Could not find a valid class-mapping for class-name '%s'. "
            "BoundingBox will be skipped",
            label_studio_class_name,
        )
        return None

    # LabelStudio has defined its rotation around the top-left point of the bounding box.
    # The RotatedBox of the MLCVZoo uses the opencv definition and rotates around the center
    # point of the bounding box. Therefore, we have to transform the LabelStudio definition
    # to the opencv definition.

    top_left_x = float(image_shape[1] * label_studio_value["x"] / 100)
    top_left_y = float(image_shape[0] * label_studio_value["y"] / 100)

    dim_x = float(image_shape[1] * label_studio_value["width"] / 100)
    dim_y = float(image_shape[0] * label_studio_value["height"] / 100)

    rotation = label_studio_value["rotation"]

    content = label_studio_value.get("text", [""])[0]

    if rotation > 0.0:
        top_left: Point2f = (top_left_x, top_left_y)
        top_right: Point2f = (top_left_x + dim_x, top_left_y)
        bottom_right: Point2f = (top_left_x + dim_x, top_left_y + dim_y)
        bottom_left: Point2f = (top_left_x, top_left_y + dim_y)

        top_right = rotate_point(top_right, top_left, rotation)
        bottom_right = rotate_point(bottom_right, top_left, rotation)
        bottom_left = rotate_point(bottom_left, top_left, rotation)

        box = np.asarray(  # type: ignore[var-annotated]
            [top_left, top_right, bottom_right, bottom_left], dtype=np.int_
        )
        (center_x, center_y), (w, h), angle = cv2.minAreaRect(box)

        return BoundingBox(
            class_identifier=ClassIdentifier(
                class_id=class_id,
                class_name=class_name,
            ),
            box=Box.init_format_based(
                box_list=(float(center_x), float(center_y), float(w), float(h)),
                box_format=ObjectDetectionBBoxFormats.CXCYWH,
                angle=angle,
            ),
            score=1.0,
            difficult=False,
            occluded=False,
            content=content,
            background=False,
        )

    return BoundingBox(
        class_identifier=ClassIdentifier(
            class_id=class_id,
            class_name=class_name,
        ),
        box=Box.init_format_based(
            box_list=(top_left_x, top_left_y, dim_x, dim_y),
            box_format=ObjectDetectionBBoxFormats.XYWH,
            angle=0.0,
        ),
        score=1.0,
        difficult=False,
        occluded=False,
        content=content,
        background=False,
    )


def parse_segmentation(
    mapper: AnnotationClassMapper, label_studio_value: Dict[str, Any], image_shape: Tuple[int, int]
) -> Optional[Segmentation]:
    """
    Try to parse a single Segmentation object from a given label studio result value (one
    entry of a result in the annotation file).

    Args:
        mapper: The mapper that should be used to translate class names
        label_studio_value: The input value for parsing the BoundingBox
        image_shape: The shape of the related image

    Returns:
        Optional[Segmentation]: The parsed Segmentation object
    """
    polygonlabels = label_studio_value.get("polygonlabels", None)
    labels = label_studio_value.get("labels", None)
    if polygonlabels:
        label_studio_class_name = polygonlabels[0]
    elif labels:
        label_studio_class_name = labels[0]
    else:
        # We don't except annotations without a class
        return None

    try:
        class_name = mapper.map_annotation_class_name_to_model_class_name(
            class_name=label_studio_class_name
        )
        class_id = mapper.map_annotation_class_name_to_model_class_id(
            class_name=label_studio_class_name
        )
    except ClassMappingNotFoundError:
        logger.debug(
            "Could not find a valid class-mapping for class-name '%s'. "
            "Segmentation will be skipped",
            label_studio_class_name,
        )
        return None

    return Segmentation(
        class_identifier=ClassIdentifier(
            class_id=class_id,
            class_name=class_name,
        ),
        polygon=[
            [x / 100 * image_shape[1], y / 100 * image_shape[0]]
            for (x, y) in label_studio_value["points"]
        ],
        score=1.0,
        difficult=False,
        occluded=False,
        content=label_studio_value.get("text", [""])[0],
        background=False,
    )


def parse_from_annotation(
    results: Dict[Any, Any], mapper: AnnotationClassMapper
) -> Tuple[Optional[Tuple[int, int]], List[BoundingBox], List[Segmentation]]:

    logger.warning(
        "DEPRECATED: parse_from_annotation is deprecated and " "will be removed in future versions"
    )

    label_studio_annotation_data: LabelStudioAnnotationData = parse_annotation_data(
        results=results, mapper=mapper
    )

    return (
        label_studio_annotation_data.image_shape,
        label_studio_annotation_data.bounding_boxes,
        label_studio_annotation_data.segmentations,
    )


def parse_annotation_data(
    results: Dict[Any, Any], mapper: AnnotationClassMapper
) -> LabelStudioAnnotationData:
    """
    Parse all bounding-boxes and segmentations that are present in the results
    of a label studio annotation. If available parse the image-shape that should
    be associated with this annotation.

    Args:
        results: The input results for the parsing
        mapper: The mapper that should be used to translate class names

    Returns:
        Tuple[Optional[Tuple[int, int]], List[Classification], List[BoundingBox], List[Segmentation]]:
        The parsed data
    """
    image_shape: Optional[Tuple[int, int]] = None
    classifications: List[Classification] = []
    bounding_boxes: List[BoundingBox] = []
    segmentations: List[Segmentation] = []

    for result in results:
        result_types = result["type"]
        result_value = result["value"]

        try:
            image_shape = (int(result["original_height"]), int(result["original_width"]))
        except KeyError:
            # Classification tasks don't contain the entries
            # "original_height" and "original_width"
            image_shape = None

        if "choices" in result_types:
            classification: Optional[Classification] = parse_classification(
                mapper=mapper,
                label_studio_value=result_value,
            )
            if classification is not None:
                classifications.append(classification)

        if (
            "rectangle" in result_types or "rectanglelabels" in result_types
        ) and image_shape is not None:
            bounding_box: Optional[BoundingBox] = parse_bounding_box(
                mapper=mapper,
                label_studio_value=result_value,
                image_shape=image_shape,
            )
            if bounding_box is not None:
                bounding_boxes.append(bounding_box)
        elif (
            "polygon" in result_types or "polygonlabels" in result_types
        ) and image_shape is not None:
            segmentation: Optional[Segmentation] = parse_segmentation(
                mapper=mapper,
                label_studio_value=result_value,
                image_shape=image_shape,
            )
            if segmentation is not None:
                segmentations.append(segmentation)

    return LabelStudioAnnotationData(image_shape, classifications, bounding_boxes, segmentations)


class LabelStudioAnnotationParser(AnnotationParser):
    """
    Parser for the Label Studio json format of a json file that is
    generated by Label Studio per image

    Support data types:
    - BoundingBox
    """

    def __init__(
        self,
        mapper: AnnotationClassMapper,
        label_studio_input_data: List[AnnotationHandlerLabelStudioInputDataConfig],
    ):
        AnnotationParser.__init__(self, mapper=mapper)

        self.label_studio_input_data = label_studio_input_data

    def parse(self) -> List[BaseAnnotation]:
        annotations: List[BaseAnnotation] = []

        for dataset_count, input_data in enumerate(self.label_studio_input_data):
            image_file_dict = get_image_file_dict(
                input_dir=input_data.input_image_dir, image_format=input_data.image_format
            )

            label_studio_annotation_paths = [
                str(p) for p in Path(input_data.input_annotation_dir).glob("**/*") if p.is_file()
            ]

            for label_studio_annotation_path in label_studio_annotation_paths:
                label_studio_dict = get_label_studio_annotation_dict(
                    label_studio_annotation_path=label_studio_annotation_path
                )

                if label_studio_dict is None:
                    logger.warning(
                        "Label Studio annotation content could not be loaded from: '%s', "
                        "Annotation will be skipped." % label_studio_annotation_path
                    )
                    continue

                if not isinstance(label_studio_dict, dict):
                    logger.warning(
                        "Loaded Label Studio annotation content is not of type dict, file: '%s', "
                        "Annotation will be skipped." % label_studio_annotation_path
                    )
                    continue

                try:
                    label_studio_image_name = os.path.basename(
                        label_studio_dict["task"]["data"]["image"]
                    )

                    if label_studio_image_name in image_file_dict:
                        image_path = image_file_dict[label_studio_image_name]
                    else:
                        logger.warning(
                            "Label Studio annotation contains non-existing image: '%s', "
                            "Annotation will be skipped." % label_studio_image_name
                        )
                        continue

                    label_studio_annotation_data: LabelStudioAnnotationData = (
                        parse_annotation_data(
                            results=label_studio_dict["result"],
                            mapper=self.mapper,
                        )
                    )

                    if label_studio_annotation_data.image_shape is None:
                        _image_shape = FrameShape(*cv2.imread(image_path).shape)
                        label_studio_annotation_data.image_shape = (
                            _image_shape.height,
                            _image_shape.width,
                        )

                    annotations.append(
                        BaseAnnotation(
                            image_path=image_path,
                            annotation_path=label_studio_annotation_path,
                            image_shape=label_studio_annotation_data.image_shape,
                            classifications=label_studio_annotation_data.classifications,
                            bounding_boxes=label_studio_annotation_data.bounding_boxes,
                            segmentations=label_studio_annotation_data.segmentations,
                            image_dir=os.path.dirname(image_path),
                            annotation_dir=os.path.dirname(label_studio_annotation_path),
                            # TODO: How to set this?
                            replacement_string="",
                        )
                    )
                except (ValueError, ForbiddenClassError, KeyError) as error:
                    logger.exception(
                        msg="Could not parse annotation file '%s'. Annotation will be skipped."
                        % label_studio_annotation_path,
                        exc_info=error,
                    )

        return annotations
