"""
General utilities
"""
import numpy as np
import cv2

from itertools import combinations


def filter_small_boxes(boxes, *args, min_area=0.0):
    """Remove bounding boxes smaller than threshold

    Args:
        boxes (np.ndarray): bounding boxes
        *args (np.ndarray): list of other arrays to filter alongside bounding boxes
        min_area (float, optional): minimum bounding box size to keep. Defaults to 0.0.

    Returns:
        list, ...: filtered bounding boxes and other arguments
    """
    # calculate area and construct a mask
    area = np.product(boxes[:, 2:], axis=1)
    mask = area > min_area
    return boxes[mask, :], *[arg[mask, ...] for arg in args]


def boxes_to_square(boxes):
    """Convert bounding boxes to square

    Args:
        boxes (np.ndarray): input bounding boxes

    Returns:
        np.ndarray: converted boxes
    """
    xy = boxes[..., :2]
    wh = boxes[..., 2:]
    # final square side is maximum of width and height
    side = np.max(wh, axis=1, keepdims=True)
    transformed_boxes = np.concatenate((xy, side, side), axis=-1)
    return transformed_boxes


def boxes_xywh2xyxy(boxes):
    """Convert bounding boxes of [x_center, y_center, width, height] format to [x_min, y_min, x_max, y_max] format

    Args:
        boxes (np.ndarray): input boxes

    Returns:
        np.ndarray: converted boxes
    """
    boxes_x1y1 = boxes[..., :2] - boxes[..., 2:] / 2
    boxes_x2y2 = boxes[..., :2] + boxes[..., 2:] / 2
    boxes =  np.concatenate((boxes_x1y1, boxes_x2y2), axis=-1)
    return boxes


def boxes_xyxy2xywh(boxes):
    """Convert bounding boxes of [x_min, y_min, x_max, y_max] format to [x_center, y_center, width, height] format

    Args:
        boxes (np.ndarray): input boxes

    Returns:
        np.ndarray: converted boxes
    """
    boxes_wh = boxes[..., 2:] - boxes[..., :2]
    boxes_xy = boxes[..., :2] + boxes_wh / 2
    boxes = np.concatenate((boxes_xy, boxes_wh), axis=-1)
    return boxes


def boxes_iou(boxes1, boxes2, center_format=True, scale=1.0):
    """Calculate intersection over union between two sets of bounding boxes

    Args:
        boxes1 (np.ndarray): _description_
        boxes2 (np.ndarray): _description_
        center_format (bool, optional): indicator if bounding boxes are of [x_center, y_center, width, height]
                                        format, if False - expects [x_min, y_min, x_max, y_max] format. Defaults to True.
        scale (float, optional): bounding box scale ratio. Defaults to 1.0.

    Returns:
        np.ndarray: IOU results
    """
    # convert to [x_center, y_center, width, height] format if needed
    if not center_format:
        boxes1 = boxes_xyxy2xywh(boxes1)
        boxes2 = boxes_xyxy2xywh(boxes2)
    # scale bboxes and convert to [x_min, y_min, x_max, y_max]
    boxes1[..., 2:] = scale * boxes1[..., 2:]
    boxes2[..., 2:] = scale * boxes2[..., 2:]
    boxes1 = boxes_xywh2xyxy(boxes1)
    boxes1_x1, boxes1_y1, boxes1_x2, boxes1_y2 = boxes1[..., 0], boxes1[..., 1], boxes1[..., 2], boxes1[..., 3]
    boxes2 = boxes_xywh2xyxy(boxes2)
    boxes2_x1, boxes2_y1, boxes2_x2, boxes2_y2 = boxes2[..., 0], boxes2[..., 1], boxes2[..., 2], boxes2[..., 3]
    # intersection
    int_w = np.maximum(np.minimum(boxes1_x2, boxes2_x2) - np.maximum(boxes1_x1, boxes2_x1), 0)
    int_h = np.maximum(np.minimum(boxes1_y2, boxes2_y2) - np.maximum(boxes1_y1, boxes2_y1), 0)
    int_area = int_w * int_h
    # union
    box_1_area = (boxes1_x2 - boxes1_x1) * (boxes1_y2 - boxes1_y1)
    box_2_area = (boxes2_x2 - boxes2_x1) * (boxes2_y2 - boxes2_y1)
    result = int_area / (box_1_area + box_2_area - int_area)
    return result


def merge_boxes(boxes, scores, classes, min_iou=1e-2):
    """Merge boxes with intersection larger than threshold

    Args:
        boxes (np.ndarray): object bounding boxes of [x_center, y_center, width, height] format
        scores (np.array): object scores
        classes (np.array): object classes
        min_iou (float, optional): merge intersection threshold. Defaults to 1e-2.

    Returns:
        np.ndarray, np.array, np.array: output boxes, scores and classes
    """
    # if no objects are passed
    if boxes.shape[0] == 0:
        return boxes, scores, classes
    boxes_out, scores_out, classes_out = list(), list(), list()
    # for every class independently
    for cls in np.unique(classes):
        # boxes, scores or current class
        boxes_cls = boxes[classes == cls, :]
        scores_cls = scores[classes == cls]
        # merge boxes and scores and save results
        boxes_merged, scores_merged = merge_boxes_single_class(boxes_cls, scores_cls, min_iou=min_iou)
        boxes_out.append(boxes_merged)
        scores_out.append(scores_merged)
        classes_out.append(np.zeros_like(scores_merged, dtype=np.int32) + cls)
    # concat results to arrays
    boxes_out = np.concatenate(boxes_out, axis=0)
    scores_out = np.concatenate(scores_out, axis=0)
    classes_out = np.concatenate(classes_out, axis=0)
    return boxes_out, scores_out, classes_out


def merge_boxes_single_class(boxes, scores, min_iou=1e-2):
    """Merge boxes of the same class with intersection larger than threshold

    Args:
        boxes (np.ndarray): _description_
        scores (np.array): _description_
        min_iou (float, optional): merge intersection threshold. Defaults to 1e-2.

    Returns:
        _type_: np.ndarray, np.array: output boxes and scores
    """
    if boxes.shape[0] < 2:
        return boxes, scores
    # transform boxes to (x1, y1, x2, y2) format
    boxes = boxes_xywh2xyxy(boxes)
    # compute iou between all combinations of boxes
    indices = np.array(list(combinations(range(boxes.shape[0]), 2)), dtype=np.int32)
    boxes1, boxes2 = boxes[indices[:, 0], :], boxes[indices[:, 1], :]
    ious = boxes_iou(boxes1, boxes2, center_format=False)
    # group overlapping boxes, assigned group ids are stored in merge_group, -1 signalizes ungrouped
    merge_group = np.zeros(boxes.shape[0], dtype=np.int32) - 1
    # group counter
    group = 0
    for idx, iou in zip(indices, ious):
        # don't merge boxes if iou is lower than threshold
        if iou < min_iou:
            continue
        # both boxes are not grouped yet
        if merge_group[idx[0]] == -1 and merge_group[idx[1]] == -1:
            merge_group[idx[0]] = group
            merge_group[idx[1]] = group
            group += 1
        # both boxes are grouped
        elif merge_group[idx[0]] != -1 and merge_group[idx[1]] != -1:
            # unify overlapping groups
            group_override = merge_group[idx[1]]
            unify_group = merge_group[idx[0]]
            merge_group = [unify_group if mg == group_override else mg for mg in merge_group]
            merge_group[idx[0]] = unify_group
            merge_group[idx[1]] = unify_group
        # one box is grouped, the other one is not
        else:
            curr_group = merge_group[idx[0]] if merge_group[idx[1]] == -1 else merge_group[idx[1]]
            merge_group[idx[0]] = curr_group
            merge_group[idx[1]] = curr_group
    boxes_out = list()
    scores_out = list()
    # add non-overlapping boxes to output list
    for box, score in zip(boxes[merge_group == -1, :], scores[merge_group == -1]):
        boxes_out.append(box)
        scores_out.append(score)
    # merge boxes of the same group and add to output list
    for group in np.unique(merge_group):
        if group == -1:
            continue
        boxes_to_merge = boxes[merge_group == group, :]
        x1y1 = np.min(boxes_to_merge[:, :2], axis=0)
        x2y2 = np.max(boxes_to_merge[:, 2:], axis=0)
        xyxy = np.concatenate((x1y1, x2y2), axis=0)
        boxes_out.append(xyxy)
        scores_to_merge = scores[merge_group == group]
        score = np.mean(scores_to_merge)
        scores_out.append(score)
    boxes_out = np.array(boxes_out)
    scores_out = np.array(scores_out)
    # transform boxes back to (x, y, w, h) format
    boxes_out = boxes_xyxy2xywh(boxes_out)
    return boxes_out, scores_out


def draw_tracks(img, boxes, keypoints, ages, hand_skeleton=None, draw_text=True, draw_boxes=True):
    """Draw tracker outputs onto canvas image

    Args:
        img (np.ndarray): BGR image to draw on
        boxes (list): list of hand bounding boxes and face bounding boxes
        keypoints (list): list of hand keypoints and face keypoints
        ages (list): ages of tracked objects
        hand_skeleton (list, optional): list of hand parental keypoints representing hand skeleton. Defaults to None.
        draw_text (bool, optional): wether to draw track age on top of bounding boxes. Defaults to True.
        draw_boxes (bool, optional): wether to draw bounding boxes. Defaults to True.

    Returns:
        np.ndarray: draw on image
    """
    # draw boxes
    if draw_boxes:
        # draw hand boxes and then face boxes
        for cls, (boxes_cls, hits_cls) in enumerate(zip(boxes, ages)):
            hits_cls = hits_cls.astype(np.int32)
            # color based on class
            color = (191, 111, 39) if cls == 0 else (71, 58, 240)
            # draw boxes
            for box, age in zip(boxes_cls, hits_cls):
                x, y, w, h = box
                xy1 = (int(x-w/2), int(y-h/2))
                xy2 = (int(x+w/2), int(y+h/2))
                cv2.rectangle(img, xy1, xy2, color, 2)
                if draw_text:
                    cv2.putText(img, f'{age}', (xy1[0], xy1[1] - 10), cv2.FONT_HERSHEY_DUPLEX, 1, color, 2)
    # draw hand keypoints
    hand_keypoints = keypoints[0].astype(np.int32)
    for kpts in hand_keypoints:
        for i, pnt in enumerate(kpts):
            # draw skeleton if passed
            if hand_skeleton is not None and hand_skeleton[i] != -1:
                img = cv2.line(img, tuple(kpts[hand_skeleton[i], :2]), tuple(pnt[:2]), (202, 183, 42), 1)
            # draw keypoints
            img = cv2.circle(img, tuple(pnt[:2]), 3, (202, 183, 42), -1)
    # draw face keypoints
    face_keypoints = keypoints[1].astype(np.int32)
    for kpts in face_keypoints:
        for pnt in kpts:
            # draw keypoints
            img = cv2.circle(img, tuple(pnt), 1, (209, 96, 61), -1)
    return img
