# coding=utf-8
# Copyright 2022 The Deeplab2 Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""COCO-style instance segmentation evaluation metrics.

Implements a Keras interface to COCO API.
COCO API: github.com/cocodataset/cocoapi/
"""
from typing import Any, Collection, Mapping, Optional

from absl import logging
import numpy as np
from pycocotools.coco import COCO
from pycocotools.cocoeval import COCOeval
import tensorflow as tf

from deeplab2.utils import coco_tools
from deeplab2.utils import panoptic_instances


def _unwrap_segmentation(seg):
  return {
      'size': list(seg['size']),
      'counts': seg['counts'],
  }


_ANNOTATION_CONVERSION = {
    'bbox': list,
    'segmentation': _unwrap_segmentation,
}


def _unwrap_annotation(ann: Mapping[str, Any]) -> Mapping[str, Any]:
  """Unwraps the objects in an COCO-style annotation dictionary.

  Logic within the Keras metric class wraps the objects within the ground-truth
  and detection annotations in ListWrapper and DictWrapper classes. On the other
  hand, the COCO API does strict type checking as part of determining which
  branch to use in comparing detections and segmentations. We therefore have
  to coerce the types from the wrapper to the built-in types that COCO is
  expecting.

  Args:
    ann: A COCO-style annotation dictionary that may contain ListWrapper and
      DictWrapper objects.

  Returns:
    The same annotation information, but with wrappers reduced to built-in
    types.
  """
  unwrapped_ann = {}
  for k in ann:
    if k in _ANNOTATION_CONVERSION:
      unwrapped_ann[k] = _ANNOTATION_CONVERSION[k](ann[k])
    else:
      unwrapped_ann[k] = ann[k]
  return unwrapped_ann


class InstanceAveragePrecision(tf.keras.metrics.Metric):
  """COCO evaluation metric class."""

  def __init__(self, name: str = 'instance_ap', **kwargs):
    """Constructs COCO evaluation class."""
    super(InstanceAveragePrecision, self).__init__(name=name, **kwargs)
    self.reset_states()

  def reset_states(self) -> None:
    """Reset COCO API object."""
    self.detections = []
    self.dataset = {
        'images': [],
        'annotations': [],
        'categories': []
    }
    self.image_id = 1
    self.next_groundtruth_annotation_id = 1