#!D:/IDE/anaconda3/envs python
# -*- coding: UTF-8 -*-
"""
 @Author: Administrator
 @FileName：YOLODetectorModel.py\n
 @DateTime: 2025/7/4 13:13\n
 @Author  ：xsh
"""
import      cv2
import numpy            as np
import onnxruntime      as ort
from PyQt5              import Qt
from AppUtils import AppUtils
from AppInterfaces.AppAlgorithmInterfaces.AlgorithmModelInterface import AlgorithmModelInterface

class YOLODetectorModel(Qt.QObject, AlgorithmModelInterface):

    def __init__(self, level: str, next_model = None, last_model = None, parent: Qt.QObject = None):
        super(YOLODetectorModel, self).__init__(parent=parent)
        self.Initialization('YOLODetectorModel', self, parent)
        self.level = level
        self.NextModel = next_model if isinstance(next_model, self.__class__) else None
        self.LastModel = last_model if isinstance(last_model, self.__class__) else None
        self.loadModel      ()

    def initValue(self):
        super(YOLODetectorModel, self).initValue()
        self.conf_threshold     = 0.5
        self.iou_threshold      = 0.45
        self.class_cont         = 0.5
        self.iou_cont           = 0.1
        # Get model width and height(YOLOv8-seg only has one input)
        self.model_height, self.model_width = 960, 960  # [x.shape for x in self.session.get_inputs()][0][-2:]

        self.dtype            = np.single
        self.session          = None
        self.level            = 'Karyotyping'
        self.NextModel        = None

    def deleValue(self):
        super(YOLODetectorModel, self).deleValue()
        if self.has_session:
            del self.session
        self.conf_threshold  = None
        self.iou_threshold   = None
        self.model_width     = None
        self.model_height    = None
        self.dtype           = None
        self.session         = None
        self.level           = None
        handler: YOLODetectorModel = self.NextModel
        model: list[YOLODetectorModel] = []
        while True:
            if not handler:
                break
            model.append(handler)
            handler = handler.next_model
        for handler in model:
            try:
                del handler
            except Exception as e:
                pass
            handler = None
        self.NextModel       = None
        self.LastModel       = None

    def loadModel(self):
        from AppUtils.AppUtils import AppUtils
        from AppCore.AppDBModels.WeightModels.PreviewClassifierModel import PreviewClassifierModel
        onnx_model      : str
        conf_threshold  : float
        iou_threshold   : float
        save_address    : str
        model_width     : int
        model_height    : int
        if not self.isAvailable:
            try:
                conf_threshold, iou_threshold, model_width, model_height, onnx_model = PreviewClassifierModel.getClassifier(name=self.level)
                self.class_cont, self.iou_cont = PreviewClassifierModel.getPreviewFilter(self.level)
                if not isinstance(conf_threshold, (int, float)) or not isinstance(iou_threshold, (float, int)):
                    self.session = None
                    return
                if not isinstance(model_width, (int, float)) or not isinstance(model_height, (float, int)):
                    self.session = None
                    return
                if not isinstance(onnx_model, str):
                    self.session = None
                    return

                try:
                    self.session = ort.InferenceSession(
                        onnx_model,
                        providers=[("CUDAExecutionProvider", {"cudnn_conv_algo_search": "DEFAULT"}),
                                   'CPUExecutionProvider']
                        if ort.get_device() == "GPU"
                        else ["CPUExecutionProvider"],
                    )

                    # Numpy dtype: support both FP32 and FP16 onnx model
                    self.dtype = np.half if self.session.get_inputs()[0].type == "tensor(float16)" else np.single
                    self.conf_threshold = conf_threshold
                    self.iou_threshold  = iou_threshold
                    self.model_height   = model_height
                    self.model_width    = model_width
                except Exception as model_ex:
                    AppUtils.print_Exception(model_ex)
                    self.session = None
                    self.dtype = np.single
                    self.conf_threshold = None
                    self.iou_threshold = None
            except Exception as e:
                AppUtils.print_Exception(e)




    def unloadModel(self):
        if self.has_session:
            del self.session
            self.session        = None
            self.dtype          = np.single
            self.conf_threshold = None
            self.iou_threshold  = None
        pass



    def process(self, level, image, *args, **kwargs):
        """
        The whole pipeline: pre-process -> inference -> post-process.

        Args:
            :param level:
            :param image: (Numpy.ndarray): original input image.
        Returns:
            boxes (Numpy.ndarray): array of obboxes (n, x_y_w_h_r_conf_cls).
        """
        data: list[int] = None
        if level == self.level:
            if not self.isAvailable:
                return data
            # Pre-process
            im, ratio, (pad_w, pad_h) = self.preprocess(image)

            # Ort inference
            preds = self.session.run(None, {self.session.get_inputs()[0].name: im})

            # Post-process
            boxes = self.postprocess(
                preds,
                image=image,
                ratio=ratio,
                pad_w=pad_w,
                pad_h=pad_h
            )

            if self.isMicronucleus():
                try:
                    if boxes is not None and len(boxes) > 1:
                        from AppCore.AppDBModels.WeightModels.PreviewClassifierModel import PreviewClassifierModel
                        self.class_cont, self.iou_cont = PreviewClassifierModel.getPreviewFilter(self.level)
                        if 0.3 < self.class_cont < 1:  # class_cont外放[0.3,1]
                            boxes = boxes[boxes[:, -1] > self.class_cont]  # 筛选类别
                        iuo_22 = self.bbox_iou_internal(boxes[:, :4])
                        if self.iou_cont > 0.9 or self.iou_cont < 0:  # iou_cont外放[0,0.9]
                            self.iou_cont = 0.1
                        del_index = np.where(iuo_22 > self.iou_cont)  # 筛选IOU    # iou_cont外放
                        if del_index[0].size > 0:
                            del_index = np.stack([del_index[0], del_index[1]])
                            del_index = set(del_index.flatten())
                            boxes = np.delete(boxes, list(del_index), axis=0)
                except Exception as e:
                    AppUtils.print_Exception(e)
                    data = None
                    return data

            if len(boxes) > 0:
                data = []
                for item in boxes:
                    data.append([int(value) for value in item])
            else:
                data = None
            return data
        elif self.has_NextModel:
            return self.NextModel.process(level, image, *args, **kwargs)
        return data

    def bbox_iou_internal(self, boxes, eps=1e-7):
        """
        计算boxes内部两两之间的IoU，返回对称矩阵

        参数:
            boxes: (N, 4)数组，格式(x_center, y_center, w, h)
            eps: 防止除零
            return_symmetric: 是否返回完整对称矩阵（False则只返回上三角）

        返回:
            iou: (N, N)对称矩阵，对角线为1.0
        """

        boxes = np.copy(boxes)
        boxes = np.asarray(boxes, dtype=float)
        if boxes.ndim == 1:
            boxes = boxes.reshape(1, 4)

        N = boxes.shape[0]
        # 1. 计算所有框的面积 (N,)
        areas = boxes[:, 2] * boxes[:, 3]

        boxes_ = np.zeros_like(boxes)

        boxes_[:, 0] = boxes[:, 0] - boxes[:, 2] / 2
        boxes_[:, 1] = boxes[:, 1] - boxes[:, 3] / 2
        boxes_[:, 2] = boxes[:, 0] + boxes[:, 2] / 2
        boxes_[:, 3] = boxes[:, 1] + boxes[:, 3] / 2

        # 2. 广播计算交集坐标 (N, N)
        #    每个框与其他所有框的xmin, ymin取最大值，xmax, ymax取最小值
        xmin = np.maximum(boxes_[:, 0, None], boxes_[:, 0])  # (N, N)
        ymin = np.maximum(boxes_[:, 1, None], boxes_[:, 1])
        xmax = np.minimum(boxes_[:, 2, None], boxes_[:, 2])
        ymax = np.minimum(boxes_[:, 3, None], boxes_[:, 3])

        # 3. 计算交集面积
        inter_w = np.maximum(0, xmax - xmin)
        inter_h = np.maximum(0, ymax - ymin)
        inter = inter_w * inter_h

        # 4. 计算并集面积 (N, N)
        union = areas[:, None] + areas - inter

        # 5. IoU矩阵
        iou_matrix = inter / (union + eps)
        mask_mat = np.triu(np.ones((N, N)), k=1)
        iou_matrix = iou_matrix * mask_mat

        return iou_matrix

    def preprocess(self, img):
        """
        Pre-processes the input image.

        Args:
            img (Numpy.ndarray): image about to be processed.

        Returns:
            img_process (Numpy.ndarray): image preprocessed for inference.
            ratio (tuple): width, height ratios in letterbox.
            pad_w (float): width padding in letterbox.
            pad_h (float): height padding in letterbox.
        """
        # Resize and pad input image using letterbox() (Borrowed from Ultralytics)
        shape = img.shape[:2]  # original image shape
        new_shape = (self.model_height, self.model_width)
        r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
        ratio = r
        new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))
        pad_w, pad_h = (new_shape[1] - new_unpad[0]) / 2, (new_shape[0] - new_unpad[1]) / 2  # wh padding
        if shape[::-1] != new_unpad:  # resize
            img = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR)
        top, bottom = int(round(pad_h - 0.1)), int(round(pad_h + 0.1))
        left, right = int(round(pad_w - 0.1)), int(round(pad_w + 0.1))
        img = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=(230, 230, 230))

        # Transforms: HWC to CHW -> BGR to RGB -> div(255) -> contiguous -> add axis(optional)
        img = np.ascontiguousarray(np.einsum("HWC->CHW", img)[::-1], dtype=self.dtype) / 255.0
        # img_min = img.min()
        # img_max = img.max()
        # img = img / 255.0   # (img - img_min) / (img_max - img_min)
        img_process = img[None] if len(img.shape) == 3 else img
        return img_process, ratio, (pad_w, pad_h)

    def postprocess(self, preds: np.ndarray, image: np.ndarray, ratio: float = 1.0, pad_w=0, pad_h=0):
        """
        Post-process the prediction.

        Args:
            preds (Numpy.ndarray): predictions come from ort.session.run().
            image (Numpy.ndarray): [h, w, c] original input image.
            ratio (tuple): width, height ratios in letterbox.
            pad_w (float): width padding in letterbox.
            pad_h (float): height padding in letterbox.

        Returns:
            boxes (List): list of bounding boxes.
            segments (List): list of segments.
            masks (np.ndarray): [N, H, W], output masks.
        """
        x = preds[0]  # Two outputs: predictions(n, x_y_w_h_obj-conf_cls-conf) and protos
        # 交换最后两个维度。
        x = x.transpose(0, 2, 1)

        if x.shape[1] < 1000:
            x = np.einsum("bcn->bnc", x)
        x = x[0]
        nc = x.shape[1] - 4

        if nc == 0:
            x = np.c_[x, np.ones(x.shape[0])]
        nc = x.shape[1] - 4
        mi = 4 + nc  # mask start index
        # Predictions filtering by conf-threshold
        # x = x[np.amax(x[..., 5:], axis=-1) > self.conf_threshold]

        # Compute conf
        # x[:, 5:] *= x[:, 4:5]  # conf = obj_conf * cls_conf

        x = x[np.amax(x[..., 4:], axis=-1) > self.conf_threshold]
        if len(x) == 0:
            return []

        x[..., 0:2] -= [pad_w, pad_h]
        x[..., 0:4] = x[..., 0:4] / ratio

        # 假设 x 是一个 NumPy 数组，形状为 (num_boxes, num_features)
        class_scores = x[:, 4:mi]  # 提取类别置信度部分

        # 获取每个检测框的最大类别置信度（保持维度）
        conf = np.max(class_scores, axis=1, keepdims=True)

        # 获取每个检测框的最大类别索引（保持维度）
        j = np.argmax(class_scores, axis=1, keepdims=True)
        # Create a new matrix which merge these(xywh, score, cls) into one

        x = np.c_[x[..., :4], conf, j]
        boxes = np.copy(x[:, :4])  # [x, y, w, h]
        c = x[:, 5:] * max(image.shape[0:2])  # 处理类别
        boxes[:, :2] += c
        # index_id = self.nms_with_iou(boxes, x[:, 4])
        # if index_id.size > 0:
        #     x2 = x[index_id, ::]
        #     return x2
        # return []

        x_ = x[cv2.dnn.NMSBoxes(boxes, x[:, 4], self.conf_threshold, self.iou_threshold)]  # 单类别
        # x_ = x[cv2.dnn.NMSBoxesBatched(x[:, :4], x[:, 4], x[:, 5].astype(np.int32), self.conf_threshold, self.iou_threshold)] # 多类别
        if x_.size > 0:
            return x_
        else:
            return []

    @classmethod
    def mask_iou(cls, mask1, mask2):
        """
        计算两个二进制掩码的交并比 (IoU)。
        :param mask1: 第一个掩码，形状为 (H, W) 的布尔数组。
        :param mask2: 第二个掩码，形状为 (H, W) 的布尔数组。
        :return: 交并比 (IoU)。
        """
        # 计算交集
        intersection = np.logical_and(mask1, mask2)
        intersection_area = np.sum(intersection)

        # 计算并集
        union = np.logical_or(mask1, mask2)
        union_area = np.sum(union)

        # 计算 IoU
        iou = intersection_area / (union_area + 1e-6)  # 防止除以零

        if intersection_area / np.sum(mask1) > 0.8 or intersection_area / np.sum(mask2) > 0.8:
            return 1
        else:
            return iou

    def nms_with_iou(self, bbox: np.ndarray, scores: np.ndarray):
        """
            批量非极大值抑制
        Args:
            :param bbox:数组，形状为 (n, x_y_w_h)
            :param scores: 置信度分数数组，形状为 (n,)
            :return: 保留的边界框索引
        """

        scores = np.array(scores)
        bbox = self.xywh2xyxy_copy(np.array(bbox))

        # 按照置信度分数降序排序
        indices = np.argsort(scores)[::-1]

        # 存储保留的边界框索引
        keep = []

        while len(indices) > 0:
            # 获取当前置信度最高的边界框
            current = indices[0]
            keep.append(current)

            if len(indices) == 1:
                break

            # 计算当前mask与其他所有mask的IoU
            ious = np.array([self.bbox_iou(bbox[current, ::], bbox[idx, ::]) for idx in indices[1:]])

            # 保留IoU小于阈值的边界框
            indices = indices[1:][ious <= self.iou_threshold]

        return np.array(keep)

    @classmethod
    def bbox_iou(cls, bbox1: np.ndarray, bbox2: np.ndarray):
        """
        计算两个边界框的IOU。
        :param bbox1: 第一个边界框，形状为 (x, y, x, y)。
        :param bbox2: 第二个边界框，形状为 (x, y, x, y)。
        :return: IOU。
        """
        # 计算交集区域的坐标
        area1 = (bbox1[0] - bbox1[2]) * (bbox1[1] - bbox1[3])
        area2 = (bbox2[0] - bbox2[2]) * (bbox2[1] - bbox2[3])

        # 计算交集的左上角和右下角坐标
        lt = np.max([bbox1[:2], bbox2[:2]], axis=0)
        rb = np.min([bbox1[2:], bbox2[2:]], axis=0)

        wh = (rb - lt).clip(min=0)  #

        # 计算交集面积
        inter_area = wh[0] * wh[1]  #

        # 计算并集面积
        union_area = area1 + area2 - inter_area

        # 计算IoU
        iou = inter_area / union_area if union_area != 0 else 0.0
        return iou

    @classmethod
    def draw_and_visualize(cls, im: np.ndarray, bboxes: np.ndarray, vis=False, save_name: str = None):
        """
        Draw and visualize results.

        Args:
            :param im           : (np.ndarray): original image, shape [h, w, c].
            :param bboxes       : (numpy.ndarray): [n, 4], n is number of bboxes.
            :param vis          : (bool): imshow using OpenCV.
            :param save_name    : (str): save image annotated.

        Returns:
            None
        """
        # Draw rectangles and polygons
        im_canvas = im.copy()
        box_data = cls.xywh_to_corner_points(bboxes[:, :4])
        count = 0
        for box, conf, cls_ in zip(box_data, bboxes[:, 4], bboxes[:, -1]):
            count += 1

            cv2.polylines(im, np.int32([box]), True, (255, 0, 255), 2)  # white borderline
            # cv2.fillPoly(im_canvas, np.int32([segment]), (255, 0, 255))

            cv2.putText(
                im,
                f"{count}: {cls_} : {conf:.3f}",
                (int(box[0][0]), int(box[1][1] - 9)),
                cv2.FONT_HERSHEY_SIMPLEX,
                2,
                (255, 0, 0),
                5,
                cv2.LINE_AA,
            )

        # Mix image
        im = cv2.addWeighted(im_canvas, 0.3, im, 0.7, 0)

        # Show image
        if vis:
            cv2.imshow("demo", im)
            cv2.waitKey(0)
            cv2.destroyAllWindows()

        # Save image
        if save_name is not None:
            cv2.imwrite(save_name, im)

    @classmethod
    def xywhr2xyxyxyxy(cls, x):
        """
        Convert batched Oriented Bounding Boxes (OBB) from [xywh, rotation] to [xy1, xy2, xy3, xy4]. Rotation values should
        be in radians from 0 to pi/2.

        Args:
            x (numpy.ndarray): Boxes in [cx, cy, w, h, rotation] format of shape (n, 5) or (b, n, 5).

        Returns:
            (numpy.ndarray): Converted corner points of shape (n, 4, 2) or (b, n, 4, 2).
        """
        cos, sin, cat, stack = np.cos, np.sin, np.concatenate, np.stack

        ctr = x[..., :2]
        w, h, angle = (x[..., i: i + 1] for i in range(2, 5))
        cos_value, sin_value = cos(angle), sin(angle)
        vec1 = [w / 2 * cos_value, w / 2 * sin_value]
        vec2 = [-h / 2 * sin_value, h / 2 * cos_value]
        vec1 = cat(vec1, -1)
        vec2 = cat(vec2, -1)
        pt1 = ctr + vec1 + vec2
        pt2 = ctr + vec1 - vec2
        pt3 = ctr - vec1 - vec2
        pt4 = ctr - vec1 + vec2
        return stack([pt1, pt2, pt3, pt4], -2)

    @classmethod
    def xywh2xyxy_copy(cls, x):
        """
        Convert bounding box coordinates from (x, y, width, height) format to (x1, y1, x2, y2) format where (x1, y1) is the
        top-left corner and (x2, y2) is the bottom-right corner. Note: ops per 2 channels faster than per channel.

        Args:
            x (np.ndarray): The input bounding box coordinates in (x, y, width, height) format.

        Returns:
            x (np.ndarray): The bounding box coordinates in ()
        """
        x_ = np.copy(x)
        x_[..., [0, 1]] -= x_[..., [2, 3]] / 2
        x_[..., [2, 3]] += x_[..., [0, 1]]
        return x_

    def xywh_to_corner_points(self, xywh: np.ndarray):
        """
        将 xywh 格式的边界框转换为 4 个角点坐标。

        参数:
            xywh (np.ndarray): 形状为 (n, 4) 的数组，每行为 [x_center, y_center, w, h]

        返回:
            np.ndarray: 形状为 (n, 4, 2) 的数组，表示每个框的 4 个角点坐标 (x, y)
        """
        x_center, y_center, w, h = np.split(xywh, 4, axis=1)
        x1 = x_center - w / 2
        y1 = y_center - h / 2
        x2 = x_center + w / 2
        y2 = y_center - h / 2
        x3 = x_center + w / 2
        y3 = y_center + h / 2
        x4 = x_center - w / 2
        y4 = y_center + h / 2

        # 拼接成 (n, 4, 2) 的形式
        points = np.concatenate([x1, y1, x2, y2, x3, y3, x4, y4], axis=1)
        return points.reshape(-1, 4, 2)


