from typing import Optional, List

import numpy as np
from skimage.measure import label, regionprops, find_contours


def mask_to_border(mask):
    h, w = mask.shape
    border = np.zeros((h, w))
    contours = find_contours(mask, 128)
    for contour in contours:
        for c in contour:
            x = int(c[0])
            y = int(c[1])
            border[x][y] = 255

    return border


def mask_to_bbox(mask):
    """ Mask to bounding boxes """
    bboxes = []
    mask = mask_to_border(mask)
    lbl = label(mask)
    props = regionprops(lbl)
    for prop in props:
        x1 = prop.bbox[1]
        y1 = prop.bbox[0]
        x2 = prop.bbox[3]
        y2 = prop.bbox[2]
        bboxes.append([x1, y1, x2, y2])

    return bboxes


def rescale_bounding_boxes(
        original_shape: tuple, new_shape: tuple, bounding_boxes: np.ndarray
) -> np.ndarray:
    """
    Rescales bounding box coords according to new image size
    :param original_shape: (W, H)
    :param new_shape: (W, H)
    :param bounding_boxes: 2D numpy array [[x1, y1, x2, y2], ...]
    :return: scaled bbox coords 2D numpy array
    """
    original_w, original_h = original_shape
    new_w, new_h = new_shape
    bounding_boxes = bounding_boxes.astype(np.float64)
    scale_h, scale_w = new_h / original_h, new_w / original_w
    bounding_boxes[:, 0] *= scale_w
    bounding_boxes[:, 1] *= scale_h
    bounding_boxes[:, 2] *= scale_w
    bounding_boxes[:, 3] *= scale_h
    bounding_boxes = np.clip(bounding_boxes, a_min=0, a_max=None)
    bounding_boxes = bounding_boxes.astype(np.uint32)
    return bounding_boxes


def extract_bboxes(mask: np.ndarray,
                   output_size: list,
                   thr_val: float) -> Optional[list]:
    """
    Extracts bounding boxes from mask
    Args:
        mask: [H, W] mask
        output_size:
        thr_val:

    Returns:
        list of bounding boxes [[x1, y1, x2, y2], ...]
    """
    mask_h, mask_w = mask.shape
    new_h, new_w = output_size

    binary_mask = np.where(mask >= thr_val, 255, 0).astype(np.uint8)
    bboxes = mask_to_bbox(mask=binary_mask)
    if len(bboxes) == 0:  # if no bbox was found, return dummy bbox with overall score
        return None

    new_raw_bboxes = np.array(bboxes).reshape(-1, 4)
    new_raw_bboxes = rescale_bounding_boxes(original_shape=(mask_w, mask_h), new_shape=(new_w, new_h),
                                                   bounding_boxes=new_raw_bboxes[:, :4])
    new_raw_bboxes = new_raw_bboxes.tolist()
    return new_raw_bboxes


def get_overlap(pred1: np.array, pred2: np.array) -> float:
    x1 = np.max(pred1[0], pred2[0])
    y1 = np.max(pred1[1], pred2[1])
    x2 = np.min(pred1[2], pred2[2])
    y2 = np.min(pred1[3], pred2[3])
    inter = (x2 - x1).clamp(0) * (y2 - y1).clamp(0)  # clamp to zero if they don't intersect
    return inter


def if_overlaps(box: list, boxes: List[list], min_distance: int) -> bool:
    box = np.array(box)
    for b in boxes:
        x1, y1, x2, y2 = b
        b = x1 - min_distance, y1 - min_distance, x2 + min_distance, y2 + min_distance
        b = np.array(b)
        if get_overlap(box, b) > 0:
            return True
    return False