# AUTOGENERATED! DO NOT EDIT! File to edit: annotation-folder_category_adapter.ipynb (unless otherwise specified).

__all__ = ['DEFAULT_CATEGORY_FOLDER_INDEX', 'logger', 'FolderCategoryAnnotationAdapter']


# Cell
import logging
import shutil
from os.path import join, normpath, sep, basename, isfile
from ..core import assign_arg_prefix
from ..io.core import create_folder, scan_files
from .core import Annotation, AnnotationAdapter, Region, SubsetType


# Cell
DEFAULT_CATEGORY_FOLDER_INDEX = -1


# Cell
logger = logging.getLogger(__name__)


# Cell
class FolderCategoryAnnotationAdapter(AnnotationAdapter):
    """
    Adapter to read and write annotations where the folder structure represents the categories.
    """

    def __init__(self, path, categories_file_name=None, category_index=None):
        """
        Multi Label Classification Adapter to read and write annotations.
        `path`: the folder containing the annotations
        `categories_file_name`: the name of the categories file
        `category_index`: the folder index, representing the category
        """
        super().__init__(path, categories_file_name)

        self.category_index = category_index if category_index is not None else DEFAULT_CATEGORY_FOLDER_INDEX

    @classmethod
    def argparse(cls, prefix=None):
        """
        Returns the argument parser containing argument definition for command line use.
        `prefix`: a parameter prefix to set, if needed
        return: the argument parser
        """
        parser = super(FolderCategoryAnnotationAdapter, cls).argparse(prefix=prefix)
        parser.add_argument(assign_arg_prefix('--category_index', prefix),
                            dest="category_index",
                            help="The folder index, representing the category.",
                            type=int,
                            default=None)
        return parser

    def read_annotations(self, subset_type=SubsetType.TRAINVAL):
        """
        Read annotations from folder structure representing the categories.
        `subset_type`: the subset type to read
        return: the annotations as dictionary
        """
        annotations = {}
        path = join(self.path, str(subset_type))
        logger.info('Read annotations from {}'.format(path))

        file_paths = scan_files(path)

        skipped_annotations = []
        for file_path in file_paths:
            trimmed_path = self._trim_base_path(file_path, path)
            if trimmed_path not in annotations:
                annotations[trimmed_path] = Annotation(annotation_id=trimmed_path, file_path=file_path)
            annotation = annotations[trimmed_path]

            path_split = normpath(trimmed_path).lstrip(sep).split(sep)

            if len(path_split) <= abs(self.category_index):
                logger.warning("{}: No category folder found, skip annotation.".format(trimmed_path))
                skipped_annotations.append(file_path)
                continue

            category = path_split[self.category_index - 1]
            region = Region(labels=[category])
            annotation.regions.append(region)

        logger.info('Finished read annotations')
        logger.info('Annotations read: {}'.format(len(annotations)))
        if skipped_annotations:
            logger.info('Annotations skipped: {}'.format(len(skipped_annotations)))
        return annotations

    def write_annotations(self, annotations, subset_type=SubsetType.TRAINVAL):
        """
        Write annotations to folder structure representing the categories.
        `annotations`: the annotations as dictionary
        `subset_type`: the subset type to write
        return: a list of written target file paths
        """
        path = join(self.path, str(subset_type))
        logger.info('Write annotations to {}'.format(path))
        skipped_annotations = []
        copied_files = []
        for annotation in annotations.values():
            if not isfile(annotation.file_path):
                logger.warning("{}: Source file not found, skip annotation.".format(annotation.file_path))
                skipped_annotations.append(annotation.file_path)
                continue

            skipped_labels = []
            annotation_labels = annotation.labels()
            for label in annotation_labels:
                category_folder = create_folder(join(path, label))
                target_file = join(category_folder, basename(annotation.file_path))
                if isfile(target_file):
                    logger.warning("{}: Target file already exist, skip label {}.".format(annotation.file_path, label))
                    skipped_labels.append(label)
                    continue
                # copy the file
                shutil.copy2(annotation.file_path, target_file)
                copied_files.append(target_file)
            if len(skipped_labels) == len(annotation_labels):
                logger.warning("{}: All labels skipped, skip annotation.".format(annotation.file_path))
                skipped_annotations.append(annotation.file_path)

        logger.info('Finished write annotations')
        logger.info('Annotations written: {}'.format(len(annotations) - len(skipped_annotations)))
        if skipped_annotations:
            logger.info('Annotations skipped: {}'.format(len(skipped_annotations)))
        return copied_files

    @classmethod
    def _trim_base_path(cls, file_path, base_path):
        """
        Trims the base path from a file path.
        `file_path`: the file path to trim from
        `base_path`: the base path to trim
        return: the trimmed file path
        """
        if file_path.startswith(base_path):
            file_path = file_path[len(base_path):]
        return file_path
