import tensorflow as tf


class mAP:
    """Mean Average Precision object detection metric
    """
    def __init__(self, iou_threshold=0.5, classes=80, inter_points=11, ignore_area_threshold=0, single_class_labels=False):
        """
        Args:
            iou_threshold (float, optional): mAP iou threshold. Defaults to 0.5.
            classes (int, optional): number of classes. Defaults to 80.
            inter_points (int, optional): number of precision-recall curve interpolation curve points. Defaults to 11.
            ignore_area_threshold (int, optional): ignore undetected labels that are smaller than threshold. Defaults to 0.
            single_class_labels (bool, optional): flag, wether only consider single class labels. Defaults to False.
        """
        self.classes = classes
        self.inter = inter_points
        self.iou_thold = tf.constant(iou_threshold, dtype=tf.float32)
        self.area_thold = tf.constant(ignore_area_threshold, dtype=tf.float32)
        self.single_class = single_class_labels

        self.true_positives = [[0] for _ in tf.range(self.classes)]
        self.precision = [[0] for _ in tf.range(self.classes)]

        self.total_true = tf.zeros((self.classes))
        self.total_positives = tf.zeros((self.classes))
    
    def update_state(self, y_true, y_pred):
        """Update metric state

        Args:
            y_true (tf.Tensor): true object labels
            y_pred (list): detector's outputs - boxes, scores, classes
        """
        # iterate over batch
        for boxes_pred, scores_pred, classes_pred, y_true_single in zip(*y_pred[:3], y_true):
            # remove zeros padding from labels
            y_true_single = y_true_single[tf.reduce_any(y_true_single != tf.constant([0, 0, 0, 0, 0], shape=(1, 5), dtype=tf.float32), axis=-1)]
            boxes_true, classes_true = tf.split(y_true_single, (4, 1), axis=-1)
            classes_true = tf.cast(classes_true, tf.int32)[:, 0]
            # filter out invalid detections (removed by nms)
            pred_valid = scores_pred != 0
            boxes_pred = tf.boolean_mask(boxes_pred, pred_valid, axis=0)
            scores_pred = tf.boolean_mask(scores_pred, pred_valid)
            classes_pred = tf.cast(tf.boolean_mask(classes_pred, pred_valid), tf.int32)
            # if single class mode is enabled, find the most frequent class in the example
            if self.single_class:
                single_class = tf.cast(tf.argmax(tf.math.bincount(classes_true)), dtype=tf.int32)
            # for every class
            for cls in tf.range(self.classes):
                # if single class mode, dont consider other classes than the most frequent one
                if self.single_class and single_class != cls:
                    continue
                # filter by classes
                pred_idx_cls = classes_pred == cls
                boxes_pred_cls = tf.boolean_mask(boxes_pred, pred_idx_cls, axis=0)
                scores_valid = tf.boolean_mask(scores_pred, pred_idx_cls)
                true_idx_cls = classes_true == cls
                boxes_true_cls = tf.boolean_mask(boxes_true, true_idx_cls, axis=0)
                # sort predictions by score
                sort_idx = tf.argsort(scores_valid, direction='DESCENDING')
                boxes_pred_cls = tf.gather(boxes_pred_cls, sort_idx, axis=0)
                # iou all combinations
                boxes1 = tf.repeat(boxes_pred_cls, tf.shape(boxes_true_cls)[0], axis=0)
                boxes2 = tf.tile(boxes_true_cls, [tf.shape(boxes_pred_cls)[0], 1])
                iou = self.iou(boxes1, boxes2)
                # reshape to (pred, true) iou matrix
                iou = tf.reshape(iou, (tf.shape(boxes_pred_cls)[0], tf.shape(boxes_true_cls)[0]))
                # indicator of taken labels
                label_taken = tf.zeros(tf.shape(boxes_true_cls)[0], dtype=tf.bool)
                # iterate over iou matrix -> yield single predicion's iou with all labels
                for single_iou in iou:
                    # update total positives counter
                    self.total_positives = tf.tensor_scatter_nd_update(self.total_positives, [[cls]], [self.total_positives[cls] + 1])
                    # sort iou vector from highest
                    best_label_idx = tf.argsort(single_iou, direction='DESCENDING')
                    # iterate over best iou indices
                    for i in best_label_idx:
                        # if label is not taken and iou is over mAP threhsold
                        if single_iou[i] >= self.iou_thold and not label_taken[i]:
                            # mark label as taken and update true positives counter
                            label_taken = tf.tensor_scatter_nd_update(label_taken, [[i]], [True])
                            self.true_positives[cls].append(self.true_positives[cls][-1] + 1)
                            break
                    else:
                        # false positive -> dont increment true positives counter
                        self.true_positives[cls].append(self.true_positives[cls][-1])
                    # calculate new precision
                    precision = self.true_positives[cls][-1] / self.total_positives[cls]
                    self.precision[cls].append(precision)
                # dont penalize recall with labels smaller than threshold
                true_labels_area = boxes_true_cls[:, 2] * boxes_true_cls[:, 3]
                true_labels = tf.boolean_mask(boxes_true_cls, tf.logical_or(true_labels_area > self.area_thold, label_taken))
                # update total labels counter
                self.total_true = tf.tensor_scatter_nd_update(self.total_true, [[cls]], [self.total_true[cls] + tf.cast(tf.shape(true_labels)[0], tf.float32)])

    def reset_state(self):
        """Reset metric's state
        """
        self.true_positives = [[0] for _ in tf.range(self.classes)]
        self.precision = [[0] for _ in tf.range(self.classes)]
        self.total_true = tf.zeros((self.classes))
        self.total_positives = tf.zeros((self.classes))

    def result(self):
        """Calculate metric results

        Returns:
            float, tf.Tensor: results of mAP and individual classes' average precisions
        """
        # calculate average precision for every class
        average_precision = tf.zeros((self.classes))
        for cls in tf.range(self.classes):
            # list of precisions and recalls
            precision = self.precision[cls]
            recall = self.true_positives[cls] / self.total_true[cls]
            # interpolation recall values
            recall_at = tf.cast(tf.linspace(0, 1, self.inter), tf.float32)
            # find precision values for recall interpolation values
            precision_at = tf.zeros((self.inter), dtype=tf.float32)
            for i, rec in zip(tf.range(self.inter), recall_at):
                # if current point is not present in recall list -> recall value not reached, precision is 0
                if not tf.reduce_any(recall >= rec):
                    update = 0
                # recall value or greater is in recall list
                else:
                    # index of the closest recall value from the list to the interpolation point
                    idx = tf.argmax(recall >= rec)
                    # maximum precision value at recall >= interpolation point
                    update = tf.reduce_max(precision[idx:])
                # update precision array
                precision_at = tf.tensor_scatter_nd_update(precision_at, [[i]], [update])
            # update average precision array
            average_precision = tf.tensor_scatter_nd_update(average_precision, [[cls]], [tf.reduce_mean(precision_at)])
        # mean of average precisions -> mAP
        return tf.reduce_mean(average_precision), average_precision

    @staticmethod
    def iou(boxes1, boxes2):
        """Intersection over union

        Args:
            boxes1 (tf.Tensor): bounding boxes in format [x_center, y_center, width, height]
            boxes2 (tf.Tensor): bounding boxes in format [x_center, y_center, width, height]

        Returns:
            tf.Tensor: IOU results between input boxes
        """
        # convert boxes to [x_min, y_min, x_max, y_max]
        boxes_1_x1, boxes_1_x2 = (tf.maximum(boxes1[..., 0] - boxes1[..., 2] / 2, 0),
                                  tf.minimum(boxes1[..., 0] + boxes1[..., 2] / 2, 1))
        boxes_1_y1, boxes_1_y2 = (tf.maximum(boxes1[..., 1] - boxes1[..., 3] / 2, 0),
                                  tf.minimum(boxes1[..., 1] + boxes1[..., 3] / 2, 1))
        boxes_2_x1, boxes_2_x2 = (tf.maximum(boxes2[..., 0] - boxes2[..., 2] / 2, 0),
                                  tf.minimum(boxes2[..., 0] + boxes2[..., 2] / 2, 1))
        boxes_2_y1, boxes_2_y2 = (tf.maximum(boxes2[..., 1] - boxes2[..., 3] / 2, 0),
                                  tf.minimum(boxes2[..., 1] + boxes2[..., 3] / 2, 1))
        # intersection
        inter_x = tf.maximum(tf.minimum(boxes_1_x2, boxes_2_x2) - tf.maximum(boxes_1_x1, boxes_2_x1), 0)
        inter_y = tf.maximum(tf.minimum(boxes_1_y2, boxes_2_y2) - tf.maximum(boxes_1_y1, boxes_2_y1), 0)
        intersection = inter_x * inter_y
        # union
        boxes_1_area = boxes1[..., 2] * boxes1[..., 3]
        boxes_2_area = boxes2[..., 2] * boxes2[..., 3]
        union = boxes_1_area + boxes_2_area - intersection
        return intersection / union
