import cv2
import numpy as np
from mindx.sdk import Tensor
from mindx.sdk import base


class YOLOV5Inference:
    def __init__(self, model_path, input_size=(640, 640), swap=(2, 0, 1),
                 score_thr=0.3, nms_thr=0.45, class_agnostic=True):
        self.model_path = model_path
        self.input_size = input_size
        self.swap = swap
        self.score_thr = score_thr
        self.nms_thr = nms_thr
        self.class_agnostic = class_agnostic

        base.mx_init()
        self.model = base.model(modelPath=self.model_path)
        if self.model is None:
            raise Exception("模型导入失败！请检查model_path。")

    def input_processing(self, image_path):
        img = cv2.imread(image_path)
        if len(img.shape) == 3:
            padded_img = np.ones((self.input_size[0], self.input_size[1], 3), dtype=np.uint8) * 114
        else:
            padded_img = np.ones(self.input_size, dtype=np.uint8) * 114

        r = min(self.input_size[0] / img.shape[0], self.input_size[1] / img.shape[1])
        resized_img = cv2.resize(
            img,
            (int(img.shape[1] * r), int(img.shape[0] * r)),
            interpolation=cv2.INTER_LINEAR,
        ).astype(np.uint8)
        padded_img[: int(img.shape[0] * r), : int(img.shape[1] * r)] = resized_img

        padded_img = padded_img.transpose(self.swap).copy()
        padded_img = np.ascontiguousarray(padded_img, dtype=np.float32)
        height, width, channels = img.shape
        return padded_img, r, height, width, channels

    def predict(self, image_path):
        img, ratio, height, width, channels = self.input_processing(image_path)
        input_tensors = img[None, :, :, :]
        input_tensors = Tensor(input_tensors)

        outputs = self.model.infer([input_tensors])[0]
        outputs.to_host()
        outputs = np.array(outputs)
        dets = self.output_processing(outputs, ratio)
        return dets

    def output_processing(self, outputs, ratio):
        outputs = outputs[0]  # [1, N, 85] -> [N, 85]
        boxes = outputs[:, 0:4]
        obj_conf = outputs[:, 4]
        class_scores = outputs[:, 5:]

        scores = obj_conf[:, None] * class_scores
        dets = multiclass_nms(boxes, scores, nms_thr=self.nms_thr, score_thr=self.score_thr,
                              class_agnostic=self.class_agnostic)

        if dets is None or len(dets) == 0:
            return np.zeros((0, 6))

        # 恢复原图尺度
        dets[:, :4] /= ratio
        return dets


def nms(boxes, scores, nms_thr):
    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
        ovr = inter / (areas[i] + areas[order[1:]] - inter)

        inds = np.where(ovr <= nms_thr)[0]
        order = order[inds + 1]

    return keep


def multiclass_nms_class_agnostic(boxes, scores, nms_thr, score_thr):
    cls_inds = scores.argmax(1)
    cls_scores = scores[np.arange(len(cls_inds)), cls_inds]
    valid_mask = cls_scores > score_thr

    if valid_mask.sum() == 0:
        return None

    valid_boxes = boxes[valid_mask]
    valid_scores = cls_scores[valid_mask]
    valid_cls_inds = cls_inds[valid_mask]
    keep = nms(valid_boxes, valid_scores, nms_thr)

    if not keep:
        return None

    dets = np.concatenate(
        [valid_boxes[keep], valid_scores[keep, None], valid_cls_inds[keep, None]], axis=1
    )
    return dets


def multiclass_nms_class_aware(boxes, scores, nms_thr, score_thr):
    final_dets = []
    num_classes = scores.shape[1]
    for cls_ind in range(num_classes):
        cls_scores = scores[:, cls_ind]
        valid_mask = cls_scores > score_thr
        if valid_mask.sum() == 0:
            continue
        valid_boxes = boxes[valid_mask]
        valid_scores = cls_scores[valid_mask]
        keep = nms(valid_boxes, valid_scores, nms_thr)
        if not keep:
            continue
        cls_inds = np.ones((len(keep), 1)) * cls_ind
        dets = np.concatenate([valid_boxes[keep], valid_scores[keep, None], cls_inds], axis=1)
        final_dets.append(dets)

    if not final_dets:
        return None
    return np.concatenate(final_dets, axis=0)


def multiclass_nms(boxes, scores, nms_thr, score_thr, class_agnostic=True):
    if class_agnostic:
        return multiclass_nms_class_agnostic(boxes, scores, nms_thr, score_thr)
    else:
        return multiclass_nms_class_aware(boxes, scores, nms_thr, score_thr)