"""Object detector utilities, used by both YOLOv5 and EfficientDet architectures
"""
import tensorflow as tf


def boxes_from_output(outputs, anchors, classes):
    """Decodes outputs from model's target grid

    Args:
        outputs (tf.Tensor): model's outputs
        anchors (tf.Tensor): model's anchor boxes
        classes (int): number of classes

    Returns:
        tf.Tensor, tf.Tensor, tf.Tensor, tf.Tensor: Decoded outputs: transformed bounding boxes to image space, objectness indicators,
                                                    class probabilities, bounding boxes in predicted format
    """
    # (batch_size, grid, grid, anchors, (x, y, w, h, obj, ...classes))
    box_xy, box_wh, objectness, class_probs = tf.split(outputs, (2, 2, 1, classes), axis=-1)
    # convert to 0~1 range
    box_xy = tf.sigmoid(box_xy)
    objectness = tf.sigmoid(objectness)
    class_probs = tf.sigmoid(class_probs)
    # bounding boxes - untransformed
    pred_box = tf.concat((box_xy, box_wh), axis=-1)
    # log-space bounding box prediction - offset to predefined anchors
    grid_size = tf.shape(outputs)[1]
    # prep grid
    grid = tf.meshgrid(tf.range(grid_size), tf.range(grid_size))
    grid = tf.expand_dims(tf.stack(grid, axis=-1), axis=2)
    # add grid offset to bbox centroids and norm to <0, 1>
    box_xy = (box_xy + tf.cast(grid, box_xy.dtype)) / tf.cast(grid_size, box_xy.dtype)
    # bbox width and height is predicted in log space and normed by anchor box size, reverse this to transform w and h to image space
    box_wh = tf.exp(box_wh) * tf.cast(anchors, box_wh.dtype)
    # concat centroids with widths and heights
    bbox = tf.concat([box_xy, box_wh], axis=-1)
    return bbox, objectness, class_probs, pred_box


def non_maximum_suppression(outputs, classes, max_detections, iou_threshold, score_threshold):
    """Apply non-max suppression to predicted bounding boxes

    Args:
        outputs (list): list od decoded outputs from all model's target grids
        classes (int): number of classes
        max_detections (int): maximum amount of detections to keep
        iou_threshold (float): NMS is applied onto detections when their IOU crosses the threshold
        score_threshold (float): minimal score of a detection to keep

    Returns:
        tf.Tensor, tf.Tensor, tf.Tensor, int: filtered bounding boxes, scores, classes and number of valid output detections
    """
    # extract bboxes, confidences and class probabilities from model's outputs
    bbox, confidence, class_probs = list(), list(), list()
    for output in outputs:
        bbox.append(tf.reshape(output[0], (tf.shape(output[0])[0], -1, tf.shape(output[0])[-1])))
        confidence.append(tf.reshape(output[1], (tf.shape(output[1])[0], -1, tf.shape(output[1])[-1])))
        class_probs.append(tf.reshape(output[2], (tf.shape(output[2])[0], -1, tf.shape(output[2])[-1])))
    # concatenate all detection together
    bbox = tf.concat(bbox, axis=1)
    confidence = tf.concat(confidence, axis=1)
    class_probs = tf.concat(class_probs, axis=1)
    # IOU non max suppression
    scores = confidence * class_probs
    # tf.image.combined_non_max_suppression expects boxes in [x_min, y_min, x_min, x_max] format
    x1y1 = bbox[:, :, :2] - bbox[:, :, 2:] / 2
    x2y2 = bbox[:, :, :2] + bbox[:, :, 2:] / 2
    bbox = tf.concat((x1y1, x2y2), axis=-1)
    boxes, scores, classes, valid_detections = tf.image.combined_non_max_suppression(
        boxes=tf.cast(tf.reshape(bbox, (tf.shape(bbox)[0], -1, 1, 4)), dtype=tf.float32),
        scores=tf.cast(tf.reshape(scores, (tf.shape(scores)[0], -1, tf.shape(scores)[-1])), dtype=tf.float32),
        max_output_size_per_class=max_detections,
        max_total_size=max_detections,
        iou_threshold=iou_threshold,
        score_threshold=score_threshold
    )
    # convert filtered boxes back to [x_center, y_center, width, height] format
    wh = boxes[:, :, 2:] - boxes[:, :, :2]
    xy = boxes[:, :, :2] + wh / 2
    boxes = tf.concat((xy, wh), axis=-1)
    return boxes, scores, classes, valid_detections


def broadcast_iou(box_1, box_2):
    """Calculate IOU between predictions and targets

    Args:
        box_1 (tf.Tensor): Predicted bounding boxes
        box_2 (tf.Tensor): Target bounding boxes

    Returns:
        tf.Tensor: IOU between input boxes
    """
    # broadcast boxes to (..., N, (x, y, w, h))
    box_1 = tf.expand_dims(box_1, -2)
    box_2 = tf.expand_dims(box_2, 0)
    new_shape = tf.broadcast_dynamic_shape(tf.shape(box_1), tf.shape(box_2))
    box_1 = tf.broadcast_to(box_1, new_shape)
    box_2 = tf.broadcast_to(box_2, new_shape)
    # convert xywh bbow format to min max format
    box_1_x1, box_1_x2 = (tf.maximum(box_1[..., 0] - box_1[..., 2] / 2, 0),
                          tf.minimum(box_1[..., 0] + box_1[..., 2] / 2, 1))
    box_1_y1, box_1_y2 = (tf.maximum(box_1[..., 1] - box_1[..., 3] / 2, 0),
                          tf.minimum(box_1[..., 1] + box_1[..., 3] / 2, 1))
    box_2_x1, box_2_x2 = (tf.maximum(box_2[..., 0] - box_2[..., 2] / 2, 0),
                          tf.minimum(box_2[..., 0] + box_2[..., 2] / 2, 1))
    box_2_y1, box_2_y2 = (tf.maximum(box_2[..., 1] - box_2[..., 3] / 2, 0),
                          tf.minimum(box_2[..., 1] + box_2[..., 3] / 2, 1))
    # intersection width, height and area
    int_w = tf.maximum(tf.minimum(box_1_x2, box_2_x2) - tf.maximum(box_1_x1, box_2_x1), 0)
    int_h = tf.maximum(tf.minimum(box_1_y2, box_2_y2) - tf.maximum(box_1_y1, box_2_y1), 0)
    int_area = int_w * int_h
    # boxes' areas
    box_1_area = box_1[..., 2] * box_1[..., 3]
    box_2_area = box_2[..., 2] * box_2[..., 3]
    return int_area / (box_1_area + box_2_area - int_area)


def DetectionLoss(anchors, classes=80, ignore_thresh=0.35):
    """Loss function for object detection model training

    Args:
        anchors (tf.Tensor): anchor boxes
        classes (int, optional): number of classes. Defaults to 80.
        ignore_thresh (float, optional): penalization threshold for prediction with greater iou with a target. Defaults to 0.35.

    Returns:
        callable: callable loss function: loss(y_true, y_pred)
    """
    def detection_loss(y_true, y_pred):
        # y_pred: (batch_size, grid, grid, anchors, (x, y, w, h, obj, ...cls))
        pred_box, pred_obj, pred_class, pred_xywh = boxes_from_output(y_pred, anchors, classes)
        pred_xy = pred_xywh[..., 0:2]
        pred_wh = pred_xywh[..., 2:4]
        # y_true: (batch_size, grid, grid, anchors, (x, y, w, h, obj, cls))
        true_box, true_obj, true_class_idx = tf.split(y_true, (4, 1, 1), axis=-1)
        true_xy = true_box[..., 0:2]
        true_wh = true_box[..., 2:4]
        # assign high weight to small boxes
        box_loss_scale = 2 - true_wh[..., 0] * true_wh[..., 1]
        # invert pred boxes equation
        grid_size = tf.shape(y_true)[1]
        grid = tf.meshgrid(tf.range(grid_size), tf.range(grid_size))
        grid = tf.expand_dims(tf.stack(grid, axis=-1), axis=2)
        true_xy = true_xy * tf.cast(grid_size, true_xy.dtype) - tf.cast(grid, true_xy.dtype)
        true_wh = tf.math.log(true_wh / anchors)
        # replace invalid values with zeros
        true_wh = tf.where(tf.math.is_inf(true_wh), tf.zeros_like(true_wh), true_wh)
        # ignore false positive when iou is over threshold
        obj_mask = tf.squeeze(true_obj, -1)
        best_iou = tf.map_fn(
            lambda x: tf.reduce_max(broadcast_iou(x[0], tf.boolean_mask(x[1], tf.cast(x[2], tf.bool))), axis=-1),
            (pred_box, true_box, obj_mask),
            fn_output_signature=tf.float32
        )
        ignore_mask = tf.cast(best_iou < ignore_thresh, obj_mask.dtype)
        # calculate losses, bboxes loss is squared error and objectness loss is binary cross-entropy
        xy_loss = obj_mask * box_loss_scale * tf.reduce_sum(tf.square(true_xy - pred_xy), axis=-1)
        wh_loss = obj_mask * box_loss_scale * tf.reduce_sum(tf.square(true_wh - pred_wh), axis=-1)
        # add weight to objectness
        obj_loss = 2 * tf.keras.losses.binary_crossentropy(true_obj, pred_obj)
        obj_loss = obj_mask * obj_loss + (1 - obj_mask) * ignore_mask * obj_loss
        # convert to one hot label and calculate binary_crossentropy
        true_class_one_hot = tf.one_hot(tf.cast(true_class_idx[..., 0], tf.int32), classes)
        class_loss = obj_mask * tf.keras.losses.binary_crossentropy(true_class_one_hot, pred_class)
        # sum along grid and anchors: (batch, gridx, gridy, anchors) -> batches
        xy_loss = tf.reduce_sum(xy_loss, axis=(1, 2, 3))
        wh_loss = tf.reduce_sum(wh_loss, axis=(1, 2, 3))
        obj_loss = tf.reduce_sum(obj_loss, axis=(1, 2, 3))
        class_loss = tf.reduce_sum(class_loss, axis=(1, 2, 3))
        return xy_loss + wh_loss + obj_loss + class_loss
    return detection_loss
