import cv2
import numpy as np


def unmold_mask(mask, bbox, image_shape):
    """
    将mask缩放到图像坐标系
    :param mask: [height, width] float型，通常较小
    :param bbox: [y1, x1, y2, x2] mask所在的box
    :param image_shape: 所需匹配的图像坐标
    :return: binary mask(图像尺寸为image_shape指定)
    """
    threshold = 0.5
    y1, x1, y2, x2 = bbox
    mask = cv2.resize(mask, (x2-x1, y2-y1))
    mask = np.where(mask >= threshold, 1, 0).astype(np.bool)

    # 将mask 放在正确的位置
    full_mask = np.zeros((int(image_shape[0]), int(image_shape[1])), dtype=np.bool)
    full_mask[y1:y2, x1:x2] = mask
    return full_mask


def get_bounding_boxes(masks):
    num_masks = masks.shape[-1]
    bounding_boxes = np.zeros([num_masks, 4], dtype=np.int32)
    for i in range(num_masks):
        m = masks[:, :, i]
        horizontal_indicies = np.where(np.any(m, axis=0))[0]
        vertical_indicies = np.where(np.any(m, axis=1))[0]
        if horizontal_indicies.shape[0]:
            x1, x2 = horizontal_indicies[[0, -1]]
            y1, y2 = vertical_indicies[[0, -1]]
            x2 += 1
            y2 += 1
        else:
            x1, x2, y1, y2 = 0, 0, 0, 0
        bounding_boxes[i] = np.array([y1, x1, y2, x2])
    return bounding_boxes


def compute_overlaps_masks(masks1, masks2):
    """
    计算两组masks的重叠率
    :param masks1: [height, width, num_instance_1]
    :param masks2: [height, width, num_instance_2]
    :return:
    """
    # 任意一个输入为空，则返回结果为空
    if masks1.shape[-1] == 0 or masks2.shape[-1] == 0:
        return np.zeros((masks1.shape[-1], masks2.shape[-2]))

    # 展平，计算面积
    masks1 = np.reshape(masks1 > .5, (-1, masks1.shape[-1])).astype(np.float32)
    masks2 = np.reshape(masks2 > .5, (-1, masks2.shape[-1])).astype(np.float32)
    area1 = np.sum(masks1, axis=0)
    area2 = np.sum(masks2, axis=0)

    # 计算IoU
    intersections = np.dot(masks1.T, masks2)
    union = area1[:, None] + area2[None, :] - intersections
    overlaps = intersections / union

    return overlaps


def compute_ratio_masks_in_rect(masks, rect):
    """
    计算每一个mask在rect中的面积比
    :param masks:
    :param rect: 未归一化
    :return:
    """
    num_mask = masks.shape[-1]
    ratios = np.zeros(num_mask, dtype=np.float32)
    for i in range(num_mask):
        area_mask = np.sum(masks[:, :, i])
        if area_mask <= 0:
            ratios[i] = 0
        area_in_rect = np.sum(masks[rect[1]: rect[3], rect[0]: rect[2], i])
        ratios[i] = area_in_rect/float(area_mask)
    return ratios


def masks_center_in_rect(masks, rect):
    """
    判断每一个mask的中心是否在rect以内
    :param masks:
    :param rect: 未归一化
    :return:
    """
    num_mask = masks.shape[-1]
    in_rect = np.zeros(num_mask, dtype=np.bool)
    for i in range(num_mask):
        inds = np.where(masks[:, :, i])
        if len(inds[0]) > 0:
            center_y = np.mean(inds[0])
            center_x = np.mean(inds[1])
            if (center_y >= rect[1]) and (center_y <= rect[3])\
                    and (center_x >= rect[0]) and (center_x <= rect[2]):
                in_rect[i] = True
    return in_rect


def clip_masks(masks, class_ids, clip_rect, in_rect=None):
    """
    如果mask的中心不在clip_rect范围内，则直接去掉该标注，如果在clip_rect范围内，则只裁取clip_rect范围内的部分
    :param masks:
    :param class_ids:
    :param clip_rect:
    :param in_rect:
    :return:
    """
    if in_rect is None:
        in_rect = masks_center_in_rect(masks, clip_rect)
    num_mask = masks.shape[-1]
    clip_rect_height = clip_rect[3] - clip_rect[1]
    clip_rect_width = clip_rect[2] - clip_rect[0]
    clipped_masks = np.zeros((clip_rect_height, clip_rect_width, num_mask), dtype=masks.dtype)
    clipped_class_ids = np.zeros_like(class_ids)
    clipped_ind = 0
    for i in range(num_mask):
        if in_rect[i]:
            clipped_masks[:, :, clipped_ind] = masks[clip_rect[1]: clip_rect[3], clip_rect[0]: clip_rect[2], i]
            clipped_class_ids[clipped_ind] = class_ids[i]
            clipped_ind += 1
    return clipped_masks, clipped_class_ids


def minimize_mask(bbox, mask, mini_shape):
    """Resize masks to a smaller version to reduce memory load.
    Mini-masks can be resized back to image scale using expand_masks()
    :param bbox: [num_gt, (y1, x1, y2, x2)] int32
    :param mask: [resize_height, resize_width, num_gt] uint8
    :param mini_shape: [mini_mask_height, mini_mask_width]
    """
    # H W C
    mini_height = mini_shape[0]
    mini_width = mini_shape[1]
    num_gt = mask.shape[-1]
    mini_mask = np.zeros((mini_height, mini_width, num_gt), dtype=np.uint8)
    for i in range(num_gt):
        m = mask[:, :, i].astype(np.uint8)
        y1, x1, y2, x2 = bbox[i][:4]
        m = m[y1:y2, x1:x2]
        if m.size == 0:
            m = np.zeros((mini_height, mini_width), dtype=np.uint8)
        else:
            m = cv2.resize(m, (mini_width, mini_height), 0, 0, interpolation=cv2.INTER_LINEAR)
        mini_mask[:, :, i] = np.around(m)
    return mini_mask