import cv2
import numpy as np


def preprocess(img_path, with_meta=False, input_size=640):
    img = cv2.imread(img_path)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    shape = img.shape[:2]
    new_shape = (input_size, input_size)
    r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
    new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))
    dw, dh = (new_shape[1] - new_unpad[0]) / 2, (new_shape[0] - new_unpad[1]) / 2
    img = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR)
    resize_h, resize_w = img.shape[:2]
    top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
    left, right = int(round(dw - 0.1)), int(round(dw + 0.1))
    img = cv2.copyMakeBorder(
        img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=(114, 114, 114)
    )
    img = img.astype(np.float32) / 255.
    img = img[None]

    if with_meta:
        img_meta = {
            'img_size': shape,
            'scale_factor': [resize_h / shape[0], resize_w / shape[1]],
            'pad_size': [dw, dh]
        }
        return img, img_meta

    return img


def postprocess(prediction, 
                img_meta, 
                conf=0.001,
                max_det=300, 
                max_nms=30000, 
                iou_thres=0.7, 
                max_wh=7680):
    # [1, 84, 8400] -> [8400, 84]
    x = prediction[0].transpose(1, 0)
    xc = np.max(x[:, 4:84], -1) > conf
    x = x[xc]

    box, cls, mask = x[:, :4], x[:, 4:4+80], x[:, 4+80:]
    i, j = np.where(cls > conf)
    x = np.concatenate((box[i], x[i, 4 + j, None], j[:, None], mask[i]), 1)
    if x.shape[0] > max_nms:  # excess boxes
        x = x[np.argsort(x[:, 4])[::-1][:max_nms]]

    x[:, :4] = xywh2xyxy(x[:, :4])
    c = x[:, 5:6] * max_wh  # classes
    boxes = x[:, :4] + c  # boxes (offset by class)

    i = nms_numpy(boxes, scores=x[:, 4], iou_threshold=iou_thres)
    i = i[:max_det]  # limit detections
    boxes, scores, labels = x[i][:, :4], x[i][:, 4], x[i][:, 5]
    boxes = scale_boxes(boxes, img_meta)

    return boxes, scores, labels


def xywh2xyxy(x):
    assert x.shape[-1] == 4, f"input shape last dimension expected 4 but input shape is {x.shape}"
    y = np.empty_like(x, dtype=np.float32)
    xy = x[..., :2]
    wh = x[..., 2:] / 2
    y[..., :2] = xy - wh 
    y[..., 2:] = xy + wh
    return y


def scale_boxes(boxes, img_meta):
    pad = img_meta['pad_size']
    scale_factor = img_meta['scale_factor']
    img_size = img_meta['img_size']
    boxes[..., [0, 2]] -= pad[0]
    boxes[..., [1, 3]] -= pad[1] 
    scale_h, scale_w = scale_factor
    scale_factor = np.array([scale_w, scale_h, scale_w, scale_h], dtype=np.float32)
    boxes[..., :4] /= scale_factor
    boxes[..., 0] = np.clip(boxes[..., 0], 0, img_size[1]) # x1
    boxes[..., 1] = np.clip(boxes[..., 1], 0, img_size[0]) # y1
    boxes[..., 2] = np.clip(boxes[..., 2], 0, img_size[1]) # x2
    boxes[..., 3] = np.clip(boxes[..., 3], 0, img_size[0]) # y2
    return boxes


def nms_numpy(boxes, scores, iou_threshold):
    if len(boxes) == 0:
        return []
    x1 = boxes[:, 0]
    y1 = boxes[:, 1]
    x2 = boxes[:, 2]
    y2 = boxes[:, 3]
    areas = (x2 - x1 + 1) * (y2 - y1 + 1)
    order = scores.argsort()[::-1]
    keep = []
    while order.size > 0:
        i = order[0] 
        keep.append(i)
        xx1 = np.maximum(x1[i], x1[order[1:]])
        yy1 = np.maximum(y1[i], y1[order[1:]])
        xx2 = np.minimum(x2[i], x2[order[1:]])
        yy2 = np.minimum(y2[i], y2[order[1:]])
        w = np.maximum(0.0, xx2 - xx1 + 1)
        h = np.maximum(0.0, yy2 - yy1 + 1)
        inter = w * h
        iou = inter / (areas[i] + areas[order[1:]] - inter)
        inds = np.where(iou <= iou_threshold)[0]
        order = order[inds + 1]
    return keep
