# 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 PascalVOC formatted annotations"""
import logging
import os
from typing import List, Tuple
from xml.etree.ElementTree import ParseError

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.exceptions import ForbiddenClassError
from mlcvzoo_base.configuration.annotation_handler_config import (
    AnnotationHandlerPASCALVOCInputDataConfig,
)
from mlcvzoo_base.configuration.structs import AnnotationFileFormats, ImageFileFormats
from mlcvzoo_base.data_preparation.annotation_builder.pascal_voc_annotation_builder import (
    PascalVOCAnnotationBuilder,
)
from mlcvzoo_base.utils.file_utils import get_file_list

logger = logging.getLogger(__name__)


class PascalVOCAnnotationParser(AnnotationParser):
    """
    Super class for defining the methods that are needed to parse a list of
    instances that are of the type BaseAnnotation.
    Each annotation format e.g. Pascal-VOC, COCO, CVAT-for-images should have
    its own child AnnotationHandler class
    """

    def __init__(
        self,
        mapper: AnnotationClassMapper,
        pascal_voc_input_data: List[AnnotationHandlerPASCALVOCInputDataConfig],
    ):
        AnnotationParser.__init__(self, mapper=mapper)

        self.pascal_voc_input_data = pascal_voc_input_data

    def parse(self) -> List[BaseAnnotation]:
        annotations: List[BaseAnnotation] = list()
        # Fill dictionary for implicit path replacement
        for dataset_count, input_data_config in enumerate(self.pascal_voc_input_data):
            replacement_string = AnnotationParser.csv_directory_replacement_string.format(
                dataset_count
            )

            image_path_list, annotation_path_list = self.__determine_data_paths(
                input_data_config=input_data_config
            )

            if input_data_config.ignore_missing_images:
                image_path_list = len(annotation_path_list) * [""]
            else:
                (
                    image_path_list,
                    annotation_path_list,
                ) = self.__ensure_annotations_match_images(
                    image_path_list=image_path_list,
                    annotation_path_list=annotation_path_list,
                )

            for image_path, annotation_path in zip(image_path_list, annotation_path_list):
                if (
                    not input_data_config.ignore_missing_images
                    and os.path.realpath(input_data_config.input_image_dir)
                    in os.path.realpath(image_path)
                    and os.path.realpath(input_data_config.input_xml_dir)
                    in os.path.realpath(annotation_path)
                ) or (
                    input_data_config.ignore_missing_images
                    and os.path.realpath(input_data_config.input_xml_dir)
                    in os.path.realpath(annotation_path)
                ):
                    try:
                        pascal_voc_builder = PascalVOCAnnotationBuilder(
                            mapper=self.mapper,
                            use_difficult=input_data_config.use_difficult,
                            use_occluded=input_data_config.use_occluded,
                            use_background=input_data_config.use_background,
                        )

                        annotation: BaseAnnotation = pascal_voc_builder.build(
                            image_path=image_path,
                            annotation_path=annotation_path,
                            image_dir=input_data_config.input_image_dir,
                            annotation_dir=input_data_config.input_xml_dir,
                            replacement_string=replacement_string,
                        )

                        annotations.append(annotation)

                    except (ForbiddenClassError, ParseError) as error:
                        logger.warning(
                            "%s: %s, annotation '%s' will be skipped"
                            % (type(error), str(error), annotation_path),
                        )
                        continue

        return annotations

    @staticmethod
    def __determine_data_paths(
        input_data_config: AnnotationHandlerPASCALVOCInputDataConfig,
    ) -> Tuple[List[str], List[str]]:
        """
        Determine all image and annotations paths from the given input-data-config

        Args:
            input_data_config: The AnnotationHandlerPASCALVOCInputDataConfig that defines
                               the input-data

        Returns:
            Tuple of image paths and annotation paths
        """

        image_paths = []
        xml_paths = []

        # CHECK SUBDIR CONFIG
        if len(input_data_config.input_sub_dirs) > 0:
            for sub_input_dir in input_data_config.input_sub_dirs:
                # READ IMAGES
                input_image_dir = os.path.join(input_data_config.input_image_dir, sub_input_dir)

                image_paths.extend(
                    get_file_list(
                        input_dir=input_image_dir,
                        search_subfolders=True,
                        file_extension=f"{input_data_config.image_format}",
                    )
                )

                # READ ANNOTATIONS
                input_xml_dir = os.path.join(input_data_config.input_xml_dir, sub_input_dir)

                xml_paths.extend(
                    get_file_list(
                        input_dir=input_xml_dir,
                        search_subfolders=True,
                        # exclude_pattern="predicted",
                        file_extension=AnnotationFileFormats.XML,
                    )
                )
        else:
            # READ IMAGES
            image_paths.extend(
                get_file_list(
                    input_dir=input_data_config.input_image_dir,
                    search_subfolders=True,
                    file_extension=input_data_config.image_format,
                )
            )

            # READ ANNOTATIONS
            xml_paths.extend(
                get_file_list(
                    input_dir=input_data_config.input_xml_dir,
                    search_subfolders=True,
                    # exclude_pattern="predicted",
                    file_extension=AnnotationFileFormats.XML,
                )
            )

        image_paths.sort()
        xml_paths.sort()

        return image_paths, xml_paths

    @staticmethod
    def __ensure_annotations_match_images(
        image_path_list: List[str], annotation_path_list: List[str]
    ) -> Tuple[List[str], List[str]]:
        """

        Args:
            image_path_list:
            annotation_path_list:

        Returns:

        """
        # remove duplicate
        image_path_list = list(dict.fromkeys(image_path_list))
        annotation_path_list = list(dict.fromkeys(annotation_path_list))

        # TODO: allow more image formats
        cleaned_img = [
            os.path.basename(image_path)
            .replace(ImageFileFormats.JPEG, "")
            .replace(ImageFileFormats.PNG, "")
            for image_path in image_path_list
        ]
        cleaned_xml = [
            os.path.basename(annotation_path)
            .replace(AnnotationFileFormats.XML, "")
            .replace("_predicted", "")
            for annotation_path in annotation_path_list
        ]

        in_both_lists = set(cleaned_xml) & set(cleaned_img)

        img_indices = [idx for (idx, x) in enumerate(cleaned_img) if x in in_both_lists]
        xml_indices = [idx for (idx, x) in enumerate(cleaned_xml) if x in in_both_lists]

        cleaned_image_path_list = [
            image_path
            for (img_index, image_path) in enumerate(image_path_list)
            if img_index in img_indices
        ]
        cleaned_annotation_path_list = [
            annotation_path
            for (annotation_index, annotation_path) in enumerate(annotation_path_list)
            if annotation_index in xml_indices
        ]

        if len(image_path_list) > len(cleaned_image_path_list):
            logger.debug(
                "Did not find an annotation-paths for image-paths: \n\t- %s\n",
                "\n\t- ".join(
                    [
                        image_path_list[index]
                        for index in range(0, len(image_path_list))
                        if index not in img_indices
                    ]
                ),
            )

        if len(annotation_path_list) > len(cleaned_annotation_path_list):
            logger.debug(
                "Did not find an image-paths for annotation-paths: \n\t- %s\n",
                "\n\t- ".join(
                    [
                        annotation_path_list[index]
                        for index in range(0, len(annotation_path_list))
                        if index not in xml_indices
                    ]
                ),
            )

        return cleaned_image_path_list, cleaned_annotation_path_list
