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

__all__ = ['DEFAULT_ANNOTATIONS_FILE', 'DEFAULT_CATEGORY_LABEL_KEY', 'CSV_FIELDNAMES', 'logger', 'VIAAnnotationAdapter']


# Cell
import json
import csv
import shutil
import logging
from os.path import join, splitext, getsize, basename, isfile
from ..core import assign_arg_prefix
from .core import Annotation, AnnotationAdapter, Region, RegionShape, parse_region_shape, SubsetType
from ..io.core import create_folder


# Cell
DEFAULT_ANNOTATIONS_FILE = 'via_region_data.json'
DEFAULT_CATEGORY_LABEL_KEY = 'category'
CSV_FIELDNAMES = ['#filename', 'file_size', 'file_attributes', 'region_count', 'region_id', 'region_shape_attributes',
                  'region_attributes']


# Cell
logger = logging.getLogger(__name__)


# Cell
class VIAAnnotationAdapter(AnnotationAdapter):
    """
    Adapter to read and write annotations in the VIA annotation.
    """

    def __init__(self, path, categories_file_name=None, annotations_file_name=None, category_label_key=None):
        """
        VIA Adapter to read and write annotations.
        `path`: the folder containing the annotations
        `categories_file_name`: the name of the categories file
        `annotations_file_name`: the name of annotations file
        `category_label_key`: the key of the category label
        """
        super().__init__(path, categories_file_name)
        if annotations_file_name is None:
            self.annotations_file_name = DEFAULT_ANNOTATIONS_FILE
        else:
            self.annotations_file_name = annotations_file_name

        if category_label_key is None:
            self.category_label_key = DEFAULT_CATEGORY_LABEL_KEY
        else:
            self.category_label_key = category_label_key

    @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(VIAAnnotationAdapter, cls).argparse(prefix=prefix)
        parser.add_argument(assign_arg_prefix('--annotations_file_name', prefix),
                            dest="annotations_file_name",
                            help="The name of annotations file.",
                            default=None)
        parser.add_argument(assign_arg_prefix('--category_label_key', prefix),
                            dest="category_label_key",
                            help="The key of the category label.",
                            default=None)

        return parser

    def read_annotations(self, subset_type=SubsetType.TRAINVAL):
        """
        Reads a VIA annotations file.
        Supports JSON and CSV file format.
        Format is inferred from the annotations_file setting.
        `subset_type`: the subset type to read
        return: the annotations as dictionary
        """
        path = join(self.path, str(subset_type))
        annotations_file_name = self._annotation_file_name_suffix_handling(subset_type)
        annotations_file_path = join(self.path, annotations_file_name)
        logger.info('Read file sources from {}'.format(path))
        logger.info('Read annotations from {}'.format(annotations_file_path))
        return self._read_annotations_v1(path, annotations_file_path)

    def _read_annotations_v1(self, path, annotations_file_path):
        """
        Reads a VIA v1 annotations file.
        Supports JSON and CSV file format.
        `path`: the path to read file sources from
        `annotation_file_path`: the path to the annotations file
        return: the annotations as dictionary
        """
        file_extension = splitext(annotations_file_path)[1]

        if file_extension.lower() == '.json':
            logger.info('Read VIA v1 annotations in JSON format')
            annotations = self._read_annotations_v1_json(path, annotations_file_path)
        elif file_extension.lower() == '.csv':
            logger.info('Read VIA v1 annotations in CSV format')
            annotations = self._read_annotations_v1_csv(path, annotations_file_path)
        else:
            message = 'Unsupported annotation format at {}'.format(annotations_file_path)
            logger.error(message)
            raise ValueError(message)

        return annotations

    def _read_annotations_v1_csv(self, path, annotations_file_path):
        """
        Reads a VIA v1 CSV annotations file.
        `path`: the path to read file sources from
        `annotation_file_path`: the path to the annotations file
        return: the annotations as dictionary
        """
        annotations = {}

        with open(annotations_file_path, newline='') as csv_file:
            reader = csv.DictReader(csv_file)

            skipped_annotations = []
            for row in reader:
                file_path = join(path, row['#filename'])
                if not isfile(file_path):
                    logger.warning("{}: Source file not found, skip annotation.".format(file_path))
                    skipped_annotations.append(file_path)
                    continue

                annotation_id = "{}{}".format(row['#filename'], row['file_size'])

                if annotation_id not in annotations:
                    annotations[annotation_id] = Annotation(annotation_id=annotation_id, file_path=file_path)

                annotation = annotations[annotation_id]

                region_shape_attributes = json.loads(row['region_shape_attributes'])
                region = self._parse_region_shape_attributes(region_shape_attributes)
                region_attributes = json.loads(row['region_attributes'])
                category = None
                if region_attributes and self.category_label_key in region_attributes:
                    category = region_attributes[self.category_label_key]
                region.labels = [category] if category else []
                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 _read_annotations_v1_json(self, path, annotations_file_path):
        """
        Reads a VIA v1 JSON annotations file.
        `path`: the path to read file sources from
        `annotation_file_path`: the path to the annotations file
        return: the annotations as dictionary
        """
        annotations = {}

        with open(annotations_file_path) as json_file:
            via_annotations = json.load(json_file)

            skipped_annotations = []
            for data in via_annotations.values():
                file_path = join(path, data['filename'])
                if not isfile(file_path):
                    logger.warning("{}: Source file not found, skip annotation.".format(file_path))
                    skipped_annotations.append(file_path)
                    continue

                annotation_id = "{}{}".format(data['filename'], data['size'])

                if annotation_id not in annotations:
                    annotations[annotation_id] = Annotation(annotation_id=annotation_id, file_path=file_path)

                annotation = annotations[annotation_id]

                for region_data in data['regions'].values():
                    region_shape_attributes = region_data['shape_attributes']
                    region = self._parse_region_shape_attributes(region_shape_attributes)
                    region_attributes = region_data['region_attributes']
                    category = None
                    if region_attributes and self.category_label_key in region_attributes:
                        category = region_attributes[self.category_label_key]
                    region.labels = [category] if category else []
                    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):
        """
        Writes a VIA annotations file and copy the corresponding source files.
        Supports JSON and CSV file format.
        Format is inferred from the annotations_file setting.
        `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))
        target_folder = create_folder(path)
        annotations_file_name = self._annotation_file_name_suffix_handling(subset_type)
        annotations_file_path = join(self.path, annotations_file_name)
        logger.info('Write file sources to {}'.format(target_folder))
        logger.info('Write annotations to {}'.format(annotations_file_path))

        return self._write_annotations_v1(path, annotations_file_path, annotations)

    def _write_annotations_v1(self, path, annotations_file_path, annotations):
        """
        Writes a VIA v1 annotations file and copy the corresponding source files.
        Supports JSON and CSV file format.
        `path`: the path to write file sources into
        `annotation_file_path`: the path to the annotations file
        `annotations`: the annotations to write
        return: a list of written source file paths
        """
        file_extension = splitext(annotations_file_path)[1]

        if file_extension.lower() == '.json':
            logger.info('Write VIA v1 annotations in JSON format')
            copied_files = self._write_annotations_v1_json(path, annotations_file_path, annotations)
        elif file_extension.lower() == '.csv':
            logger.info('Write VIA v1 annotations in CSV format')
            copied_files = self._write_annotations_v1_csv(path, annotations_file_path, annotations)
        else:
            message = 'Unsupported annotation format at {}'.format(annotations_file_path)
            logger.error(message)
            raise ValueError(message)
        return copied_files

    def _write_annotations_v1_csv(self, path, annotations_file_path, annotations):
        """
        Writes a VIA v1 CSV annotations file and copy the corresponding source files.
        `path`: the path to write file sources into
        `annotation_file_path`: the path to the annotations file
        `annotations`: the annotations to write
        return: a list of written source file paths
        """
        with open(annotations_file_path, 'w', newline='') as csv_file:
            writer = csv.DictWriter(csv_file, fieldnames=CSV_FIELDNAMES)
            writer.writeheader()

            skipped_annotations = []
            copied_files = []
            for annotation in annotations.values():
                target_file = join(path, basename(annotation.file_path))

                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
                if isfile(target_file):
                    logger.warning("{}: Target file already exist, skip annotation.".format(annotation.file_path))
                    skipped_annotations.append(annotation.file_path)
                    continue

                file_size = getsize(annotation.file_path)
                file_name = basename(annotation.file_path)
                for index, region in enumerate(annotation.regions):
                    region_shape_attributes = self._create_region_shape_attributes(region)
                    region_attributes = {
                        self.category_label_key: ' '.join(region.labels) if len(region.labels) else ''
                    }

                    writer.writerow({'#filename': file_name,
                                     'file_size': file_size,
                                     'file_attributes': '{}',
                                     'region_count': len(annotation.regions),
                                     'region_id': str(index),
                                     'region_shape_attributes': json.dumps(region_shape_attributes),
                                     'region_attributes': json.dumps(region_attributes)})
                # copy the file
                shutil.copy2(annotation.file_path, target_file)
                copied_files.append(target_file)

        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

    def _write_annotations_v1_json(self, path, annotations_file_path, annotations):
        """
        Writes a VIA v1 JSON annotations file and copy the corresponding source files.
        `path`: the path to write file sources into
        `annotation_file_path`: the path to the annotations file
        `annotations`: the annotations to write
        return: a list of written source file paths
        """
        json_annotations = {}
        skipped_annotations = []
        copied_files = []

        for annotation in annotations.values():
            target_file = join(path, basename(annotation.file_path))

            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
            if isfile(target_file):
                logger.warning("{}: Target file already exist, skip annotation.".format(annotation.file_path))
                skipped_annotations.append(annotation.file_path)
                continue

            file_size = getsize(annotation.file_path)
            file_name = basename(annotation.file_path)
            file_id = '{:s}{:d}'.format(file_name, file_size)
            regions = {}
            for index, region in enumerate(annotation.regions):
                regions[str(index)] = {
                    'shape_attributes': self._create_region_shape_attributes(region),
                    'region_attributes': {
                        self.category_label_key: ' '.join(region.labels) if len(region.labels) else ''
                    }
                }
            json_annotations[file_id] = {
                'fileref': "",
                'size': file_size,
                'filename': file_name,
                'base64_img_data': "",
                'file_attributes': '{}',
                "regions": regions
            }
            # copy the file
            shutil.copy2(annotation.file_path, target_file)
            copied_files.append(target_file)

        with open(annotations_file_path, 'w') as json_file:
            json.dump(json_annotations, json_file)

        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

    def _annotation_file_name_suffix_handling(self, subset_type):
        """
        Handle annotations file name based on the subset type.
        `subset_type`: the subset type to handle
        return: the annotations file name
        """
        file_name, ext = splitext(self.annotations_file_name)
        if subset_type in [SubsetType.TRAIN, SubsetType.VAL] and not file_name.endswith(str(subset_type)):
            suffix = "_{}".format(str(subset_type))
            return "{}{}{}".format(file_name, suffix, ext)
        return self.annotations_file_name

    @classmethod
    def _parse_region_shape_attributes(cls, region_shape_attributes):
        """
        Parse region shape attributes.
        `region_shape_attributes`: the region shape attributes as dictionary
        return: the corresponding annotation
        """
        if not region_shape_attributes:
            return Annotation()

        region_shape = parse_region_shape(region_shape_attributes['name'])
        points_x = None
        points_y = None
        radius_x = 0
        radius_y = 0
        if region_shape == RegionShape.CIRCLE:
            points_x = [region_shape_attributes['cx']]
            points_y = [region_shape_attributes['cy']]
            radius_x = region_shape_attributes['r']
            radius_y = region_shape_attributes['r']
        elif region_shape == RegionShape.ELLIPSE:
            points_x = [region_shape_attributes['cx']]
            points_y = [region_shape_attributes['cy']]
            radius_x = region_shape_attributes['rx']
            radius_y = region_shape_attributes['ry']
        elif region_shape == RegionShape.POINT:
            points_x = [region_shape_attributes['cx']]
            points_y = [region_shape_attributes['cy']]
        elif region_shape == RegionShape.POLYGON:
            points_x = region_shape_attributes['all_points_x']
            points_y = region_shape_attributes['all_points_y']
        elif region_shape == RegionShape.RECTANGLE:
            x = region_shape_attributes['x']
            y = region_shape_attributes['y']
            width = region_shape_attributes['width']
            height = region_shape_attributes['height']
            points_x = [x, x + width]
            points_y = [y, y + height]
        return Region(shape=region_shape, points_x=points_x, points_y=points_y, radius_x=radius_x, radius_y=radius_y)

    @classmethod
    def _create_region_shape_attributes(cls, region: Region):
        """
        Create region shape attributes.
        `region`: the region to create region shape attributes from
        return: the corresponding region shape attributes as dictionary
        """
        region_shape_attributes = {
            "name": str(region.shape),

        }
        c_x = region.points_x[0] if len(region.points_x) else 0
        c_y = region.points_y[0] if len(region.points_y) else 0

        if region.shape == RegionShape.CIRCLE:
            region_shape_attributes['cx'] = c_x
            region_shape_attributes['cy'] = c_y
            region_shape_attributes['r'] = max(region.radius_x, region.radius_y)
        elif region.shape == RegionShape.ELLIPSE:
            region_shape_attributes['cx'] = c_x
            region_shape_attributes['cy'] = c_y
            region_shape_attributes['rx'] = region.radius_x
            region_shape_attributes['ry'] = region.radius_y
        elif region.shape == RegionShape.POINT:
            region_shape_attributes['cx'] = c_x
            region_shape_attributes['cy'] = c_y
        elif region.shape == RegionShape.POLYGON:
            region_shape_attributes['all_points_x'] = region.points_x
            region_shape_attributes['all_points_y'] = region.points_y
        elif region.shape == RegionShape.RECTANGLE:
            region_shape_attributes['x'] = region.points_x[0]
            region_shape_attributes['y'] = region.points_y[0]
            region_shape_attributes['width'] = region.points_x[1] - region.points_x[0]
            region_shape_attributes['height'] = region.points_y[1] - region.points_y[0]
        return region_shape_attributes
