
# Copyright (C) 2019-2020 Intel Corporation
#
# SPDX-License-Identifier: MIT

from collections import OrderedDict
import logging as log
import os.path as osp

from pycocotools.coco import COCO
import pycocotools.mask as mask_utils

from datumaro.components.extractor import (SourceExtractor,
    DEFAULT_SUBSET_NAME, DatasetItem,
    AnnotationType, Label, RleMask, Points, Polygon, Bbox, Caption,
    LabelCategories, PointsCategories
)
from datumaro.util.image import Image

from .format import CocoTask, CocoPath


class _CocoExtractor(SourceExtractor):
    def __init__(self, path, task, merge_instance_polygons=False):
        assert osp.isfile(path), path

        subset = osp.splitext(osp.basename(path))[0].rsplit('_', maxsplit=1)
        subset = subset[1] if len(subset) == 2 else None
        super().__init__(subset=subset)

        rootpath = ''
        if path.endswith(osp.join(CocoPath.ANNOTATIONS_DIR, osp.basename(path))):
            rootpath = path.rsplit(CocoPath.ANNOTATIONS_DIR, maxsplit=1)[0]
        images_dir = ''
        if rootpath and osp.isdir(osp.join(rootpath, CocoPath.IMAGES_DIR)):
            images_dir = osp.join(rootpath, CocoPath.IMAGES_DIR)
            if osp.isdir(osp.join(images_dir, subset or DEFAULT_SUBSET_NAME)):
                images_dir = osp.join(images_dir, subset or DEFAULT_SUBSET_NAME)
        self._images_dir = images_dir
        self._task = task

        self._merge_instance_polygons = merge_instance_polygons

        loader = self._make_subset_loader(path)
        self._load_categories(loader)
        self._items = list(self._load_items(loader).values())

    @staticmethod
    def _make_subset_loader(path):
        # COCO API has an 'unclosed file' warning
        coco_api = COCO()
        with open(path, 'r') as f:
            import json
            dataset = json.load(f)

        coco_api.dataset = dataset
        coco_api.createIndex()
        return coco_api

    def _load_categories(self, loader):
        self._categories = {}

        if self._task in [CocoTask.instances, CocoTask.labels,
                CocoTask.person_keypoints,
                # TODO: Task.stuff, CocoTask.panoptic
                ]:
            label_categories, label_map = self._load_label_categories(loader)
            self._categories[AnnotationType.label] = label_categories
            self._label_map = label_map

        if self._task == CocoTask.person_keypoints:
            person_kp_categories = self._load_person_kp_categories(loader)
            self._categories[AnnotationType.points] = person_kp_categories

    # pylint: disable=no-self-use
    def _load_label_categories(self, loader):
        catIds = loader.getCatIds()
        cats = loader.loadCats(catIds)

        categories = LabelCategories()
        label_map = {}
        for idx, cat in enumerate(cats):
            label_map[cat['id']] = idx
            categories.add(name=cat['name'], parent=cat.get('supercategory'))

        return categories, label_map
    # pylint: enable=no-self-use

    def _load_person_kp_categories(self, loader):
        catIds = loader.getCatIds()
        cats = loader.loadCats(catIds)

        categories = PointsCategories()
        for cat in cats:
            label_id = self._label_map[cat['id']]
            categories.add(label_id=label_id,
                labels=cat['keypoints'], joints=cat['skeleton']
            )

        return categories

    def _load_items(self, loader):
        items = OrderedDict()

        for img_id in loader.getImgIds():
            image_info = loader.loadImgs(img_id)[0]
            image_path = osp.join(self._images_dir, image_info['file_name'])
            image_size = (image_info.get('height'), image_info.get('width'))
            if all(image_size):
                image_size = (int(image_size[0]), int(image_size[1]))
            else:
                image_size = None
            image = Image(path=image_path, size=image_size)

            anns = loader.getAnnIds(imgIds=img_id)
            anns = loader.loadAnns(anns)
            anns = sum((self._load_annotations(a, image_info) for a in anns), [])

            items[img_id] = DatasetItem(
                id=osp.splitext(image_info['file_name'])[0],
                subset=self._subset, image=image, annotations=anns,
                attributes={'id': img_id})

        return items

    def _get_label_id(self, ann):
        cat_id = ann.get('category_id')
        if cat_id in [0, None]:
            return None
        return self._label_map[cat_id]

    def _load_annotations(self, ann, image_info=None):
        parsed_annotations = []

        ann_id = ann.get('id')

        attributes = {}
        if 'attributes' in ann:
            try:
                attributes.update(ann['attributes'])
            except Exception as e:
                log.debug("item #%s: failed to read annotation attributes: %s",
                    image_info['id'], e)
        if 'score' in ann:
            attributes['score'] = ann['score']

        group = ann_id # make sure all tasks' annotations are merged

        if self._task in [CocoTask.instances, CocoTask.person_keypoints]:
            x, y, w, h = ann['bbox']
            label_id = self._get_label_id(ann)

            is_crowd = bool(ann['iscrowd'])
            attributes['is_crowd'] = is_crowd

            if self._task is CocoTask.person_keypoints:
                keypoints = ann['keypoints']
                points = [p for i, p in enumerate(keypoints) if i % 3 != 2]
                visibility = keypoints[2::3]
                parsed_annotations.append(
                    Points(points, visibility, label=label_id,
                        id=ann_id, attributes=attributes, group=group)
                )

            segmentation = ann.get('segmentation')
            if segmentation and segmentation != [[]]:
                rle = None

                if isinstance(segmentation, list):
                    if not self._merge_instance_polygons:
                        # polygon - a single object can consist of multiple parts
                        for polygon_points in segmentation:
                            parsed_annotations.append(Polygon(
                                points=polygon_points, label=label_id,
                                id=ann_id, attributes=attributes, group=group
                            ))
                    else:
                        # merge all parts into a single mask RLE
                        img_h = image_info['height']
                        img_w = image_info['width']
                        rles = mask_utils.frPyObjects(segmentation, img_h, img_w)
                        rle = mask_utils.merge(rles)
                elif isinstance(segmentation['counts'], list):
                    # uncompressed RLE
                    img_h = image_info['height']
                    img_w = image_info['width']
                    mask_h, mask_w = segmentation['size']
                    if img_h == mask_h and img_w == mask_w:
                        rle = mask_utils.frPyObjects(
                            [segmentation], mask_h, mask_w)[0]
                    else:
                        log.warning("item #%s: mask #%s "
                            "does not match image size: %s vs. %s. "
                            "Skipping this annotation.",
                            image_info['id'], ann_id,
                            (mask_h, mask_w), (img_h, img_w)
                        )
                else:
                    # compressed RLE
                    rle = segmentation

                if rle is not None:
                    parsed_annotations.append(RleMask(rle=rle, label=label_id,
                        id=ann_id, attributes=attributes, group=group
                    ))
            else:
                parsed_annotations.append(
                    Bbox(x, y, w, h, label=label_id,
                        id=ann_id, attributes=attributes, group=group)
                )
        elif self._task is CocoTask.labels:
            label_id = self._get_label_id(ann)
            parsed_annotations.append(
                Label(label=label_id,
                    id=ann_id, attributes=attributes, group=group)
            )
        elif self._task is CocoTask.captions:
            caption = ann['caption']
            parsed_annotations.append(
                Caption(caption,
                    id=ann_id, attributes=attributes, group=group)
            )
        else:
            raise NotImplementedError()

        return parsed_annotations

class CocoImageInfoExtractor(_CocoExtractor):
    def __init__(self, path, **kwargs):
        kwargs['task'] = CocoTask.image_info
        super().__init__(path, **kwargs)

class CocoCaptionsExtractor(_CocoExtractor):
    def __init__(self, path, **kwargs):
        kwargs['task'] = CocoTask.captions
        super().__init__(path, **kwargs)

class CocoInstancesExtractor(_CocoExtractor):
    def __init__(self, path, **kwargs):
        kwargs['task'] = CocoTask.instances
        super().__init__(path, **kwargs)

class CocoPersonKeypointsExtractor(_CocoExtractor):
    def __init__(self, path, **kwargs):
        kwargs['task'] = CocoTask.person_keypoints
        super().__init__(path, **kwargs)

class CocoLabelsExtractor(_CocoExtractor):
    def __init__(self, path, **kwargs):
        kwargs['task'] = CocoTask.labels
        super().__init__(path, **kwargs)
