# coding=utf-8
# Copyright 2022 The Google Research 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.

"""Semantic segmentation metric."""
import gin
import gin.tf

import numpy as np
import tensorflow as tf

from tf3d import standard_fields
from object_detection.utils import label_map_util


@gin.configurable
class SemanticSegmentationMetric(tf.keras.metrics.Metric):
  """Semantic segmentation mean intersection over union metric."""

  def __init__(self,
               multi_label=False,
               num_classes=None,
               label_map=None,
               label_map_path=None,
               eval_prefix='eval',
               name='semantic_segmentation_metric'):
    """Semantic segmentation mean intersection over union metric.

    Args:
      multi_label: Boolean which denotes if pixels can be assigned multiple
        labels; classes are treated separately, logit > 0 is positive
        prediction.
      num_classes: Number of classes.
      label_map: A dictionary mapping label ids to label names.
      label_map_path: path to labelmap (could be None).
      eval_prefix: Prefix for eval name; separates scalar values in Tensorboard.
      name: class name.
    """
    super(SemanticSegmentationMetric, self).__init__(name=name)
    self.multi_label = multi_label
    self.num_classes = num_classes
    if label_map:
      self.label_map = label_map
    elif label_map_path:
      self.label_map = _get_label_map(label_map_path)
    else:
      self.label_map = None

    self.eval_prefix = eval_prefix
    if self.label_map is not None:
      self.class_range = self.label_map.keys()
    elif num_classes is not None:
      self.class_range = range(num_classes)
    else:
      raise ValueError('Both num_classes and label_map are None.')
    self.true_positive_metrics = {}
    self.false_positive_metrics = {}
    self.false_negative_metrics = {}
    for c in self.class_range:
      self.true_positive_metrics[c] = tf.keras.metrics.TruePositives(
          name=('%s_true_positive_%d' % (name, c)))
      self.false_positive_metrics[c] = tf.keras.metrics.FalsePositives(
          name=('%s_false_positive_%d' % (name, c)))
      self.false_negative_metrics[c] = tf.keras.metrics.FalseNegatives(
          name=('%s_false_negative_%d' % (name, c)))

  def update_state(self, inputs, outputs):
    """Function that updates the metric state at each example.

    Args:
      inputs: A dictionary containing input tensors.
      outputs: A dictionary containing output tensors.

    Returns:
      Update op.
    """
    # Prepare logits and labels
    logits = outputs[
        standard_fields.DetectionResultFields.object_semantic_points]
    labels = inputs[standard_fields.InputDataFields.object_class_points]
    weights = inputs[standard_fields.InputDataFields.point_loss_weights]
    num_valid_points = inputs[standard_fields.InputDataFields.num_valid_points]
    if len(logits.get_shape().as_list()) == 3:
      batch_size = logits.get_shape().as_list()[0]
      logits_list = []
      labels_list = []
      weights_list = []
      for i in range(batch_size):
        num_valid_points_i = num_valid_points[i]
        logits_list.append(logits[i, 0:num_valid_points_i, :])
        labels_list.append(labels[i, 0:num_valid_points_i, :])
        weights_list.append(weights[i, 0:num_valid_points_i, :])
      logits = tf.concat(logits_list, axis=0)
      labels = tf.concat(labels_list, axis=0)
      weights = tf.concat(weights_list, axis=0)
    if self.num_classes is None:
      num_classes = logits.get_shape().as_list()[-1]
    else:
      num_classes = self.num_classes
      if num_classes != logits.get_shape().as_list()[-1]:
        raise ValueError('num_classes do not match the logits dimensions.')

    class_labels, class_predictions = _get_class_labels_and_predictions(
        labels=labels,
        logits=logits,
        num_classes=self.num_classes,
        multi_label=self.multi_label)

    update_ops = []
    for c in self.class_range:
      update_op_tp_c = self.true_positive_metrics[c].update_state(
          y_true=class_labels[c],
          y_pred=class_predictions[c],
          sample_weight=weights)
      update_ops.append(update_op_tp_c)
      update_op_fp_c = self.false_positive_metrics[c].update_state(
          y_true=class_labels[c],
          y_pred=class_predictions[c],
          sample_weight=weights)
      update_ops.append(update_op_fp_c)
      update_op_fn_c = self.false_negative_metrics[c].update_state(
          y_true=class_labels[c],
          y_pred=class_predictions[c],
          sample_weight=weights)
      update_ops.append(update_op_fn_c)
    return tf.group(update_ops)

  def result(self):
    metrics_dict = self.get_metric_dictionary()
    return metrics_dict[self.eval_prefix + '_avg/mean_iou']

  def get_metric_dictionary(self):