#!D:/IDE/anaconda3/envs python
# -*- coding: UTF-8 -*-
"""
 @Author: Administrator
 @FileName：ObbDetectorModel.py\n
 @DateTime: 2025/7/21 14:27\n
 @Author  ：xsh
"""
import      cv2
import numpy                as np
import onnxruntime          as ort
from PyQt5                  import Qt
from AppSettings.Settings   import *
from AppUtils import AppUtils
from scipy.spatial          import KDTree
from AppInterfaces.AppAlgorithmInterfaces.AlgorithmModelInterface import AlgorithmModelInterface

class ObbDetectorModel(Qt.QObject, AlgorithmModelInterface):
    def __init__(self, level: str = 'Karyotyping', parent: Qt.QObject = None):
        super(ObbDetectorModel, self).__init__(parent=parent)
        self.Initialization('ObbDetectorModel', self, parent)
        self.level = level
        self.loadModel      ()

    def __del__(self):
        self.deleValue()

    def isValid(self) -> bool:
        return self.has_level and self.has_session and self.has_dtype and self.has_conf_threshold and self.has_iou_threshold

    def initValue(self):
        super(ObbDetectorModel, self).initValue()
        self.conf_threshold   = 0.5
        self.iou_threshold    = 0.45
        # 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
        self.LastModel        = None

    def deleValue(self):
        if self.has_session:
            del self.session
        self.conf_threshold  = None
        self.iou_threshold   = None
        self.model_height    = None
        self.model_height    = None
        self.dtype           = None
        self.session         = None
        handler: ObbDetectorModel = self.NextModel
        model: list[ObbDetectorModel] = []
        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.CollectClassifierModel import CollectClassifierModel
        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 = CollectClassifierModel.getClassifier(name=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 e:
                        AppUtils.print_Exception(e)
                        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:
                level: str
                image (Numpy.ndarray): original input image.
            Returns:
                score: 大图评分
                c_point: 中心店坐标[x, y]
                count: 染色体大致数量
            """
        score           : int       = 0             #
        count           : int       = 0             #
        center_point    : list[int] = [0, 0]        #
        try:
            if level != self.level:
                if self.hasNextModel:
                    return self.NextModel.process(level, image, *args, **kwargs)
            else:
                if len(image.shape) == 2:
                    image = np.stack([image, image, image], axis=-1)
                if len(image.shape) == 3 and image.shape[-1] >= 10:  # len(im0.shape) != 4:
                    image = image.transpose(1, 2, 0)
                im, ratio, (pad_w, pad_h) = self.preprocess(image)
                preds = self.session.run(None, {self.session.get_inputs()[0].name: im})
                try:
                    obbs, iou_score = self.postprocess(
                        preds,
                        im0=image,
                        ratio=ratio,
                        pad_w=pad_w,
                        pad_h=pad_h,
                        re_oiu=True
                    )
                    if obbs is not None:
                        score, count, center_point = self.get_score(obbs, iou_score)
                        return score, count, center_point
                except Exception as e:
                    AppUtils.print_Exception(e)
        except Exception as e:
            AppUtils.print_Exception(e)
        return score, count, center_point


    def __call__(self, image: np.ndarray):
        return self.process(self.level, image)

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

            Args:
                level: str
                image (Numpy.ndarray): original input image.
            Returns:
                染色体大致数量
            """
        if level != self.level:
            if self.hasNextModel:
                return self.NextModel.img_score(level, image, *args, **kwargs)
        else:
            if len(image.shape) == 2:
                image = np.stack([image, image, image], axis=-1)
            if len(image.shape) == 3 and image.shape[-1] >= 10:  # len(im0.shape) != 4:
                image = image.transpose(1, 2, 0)

            # image_size = im0.shape[:2]
            im, ratio, (pad_w, pad_h) = self.preprocess(image)

            # Ort inference
            preds = self.session.run(None, {self.session.get_inputs()[0].name: im})
            # return 0,0
            try:
                obbs, iou_score = self.postprocess(
                    preds,
                    im0=image,
                    ratio=ratio,
                    pad_w=pad_w,
                    pad_h=pad_h,
                    re_oiu=False
                )
                if obbs is not None:
                    return int(obbs.shape[0])
            except Exception as e:
                AppUtils.print_Exception(e)
        return 0


    def probiou(self, obb1: np.ndarray, obb2: np.ndarray, eps=1e-7):
        """
        Calculate the prob IoU between oriented bounding boxes, https://arxiv.org/pdf/2106.06072v1.pdf.

        Args:
            obb1 (np.ndarray): A tensor of shape (N, 5) representing ground truth obbs, with xywhr format.
            obb2 (np.ndarray): A tensor of shape (M, 5) representing predicted obbs, with xywhr format.
            eps (float, optional): A small value to avoid division by zero. Defaults to 1e-7.

        Returns:
            (torch.Tensor): A tensor of shape (N, M) representing obb similarities.
        """
        # obb1 = torch.from_numpy(obb1) if isinstance(obb1, np.ndarray) else obb1
        # obb2 = torch.from_numpy(obb2) if isinstance(obb2, np.ndarray) else obb2
        if len(obb1.shape) == 2:
            x1, y1 = np.split(obb1[..., :2], 2, axis=-1)  # [n,1]
            x2, y2 = (np.squeeze(x, axis=-1)[None] for x in np.split(obb2[..., :2], 2, axis=-1))  # [1,n]
            a1, b1, c1 = self._get_covariance_matrix(obb1)
            a2, b2, c2 = (x.squeeze(-1)[None] for x in self._get_covariance_matrix(obb2))
        else:
            x1, y1 = obb1[0], obb1[1]
            x2, y2 = obb2[0], obb2[1]
            a1, b1, c1 = self._get_covariance_matrix(obb1)
            a2, b2, c2 = self._get_covariance_matrix(obb2)

        t1 = (((a1 + a2) * np.power((y1 - y2), 2) + (b1 + b2) * np.power((x1 - x2), 2)) / (
                (a1 + a2) * (b1 + b2) - np.power((c1 + c2), 2) + eps)) * 0.25
        t2 = (((c1 + c2) * (x2 - x1) * (y1 - y2)) / ((a1 + a2) * (b1 + b2) - np.power((c1 + c2), 2) + eps)) * 0.5
        t3 = np.log(((a1 + a2) * (b1 + b2) - np.power((c1 + c2), 2)) /
                    (4 * np.sqrt(np.clip((a1 * b1 - np.power(c1, 2)), a_min=0, a_max=None) *
                                 np.clip((a2 * b2 - np.power(c2, 2)), a_min=0, a_max=None)) + eps) + eps) * 0.5
        bd = np.clip((t1 + t2 + t3), eps, 100.0)
        hd = np.sqrt(1.0 - np.exp(-bd) + eps)
        return 1 - hd

    def _get_covariance_matrix(self, boxes: np.ndarray):
        """
        Generating covariance matrix from obbs.

        Args:
            boxes (np.ndarray): A tensor of shape (N, 5) or (5,) representing rotated bounding boxes, with xywhr format.

        Returns:
            (np.ndarray / float): Covariance matrices corresponding to original rotated bounding boxes.
        """
        # Gaussian bounding boxes, ignore the center points (the first two columns) because they are not needed here.
        if len(boxes.shape) == 2:
            gbbs = np.concatenate((np.power(boxes[:, 2:4], 2) / 12, boxes[:, 4:]), axis=-1)
            a, b, c = np.split(gbbs, 3, axis=-1)
        else:
            gbbs = np.concatenate((np.power(boxes[2:4], 2) / 12, boxes[4:]), axis=-1)
            a, b, c = gbbs[0], gbbs[1], gbbs[2]

        cos = np.cos(c)
        sin = np.sin(c)
        cos2 = np.power(cos, 2)
        sin2 = np.power(sin, 2)
        return a * cos2 + b * sin2, a * sin2 + b * cos2, (a - b) * cos * sin

    def nms_with_iou(self, obbox: np.ndarray, scores: np.ndarray, iou_thres=None, re_iou=False):
        """
            批量非极大值抑制
        Args:
            obbox:数组，形状为 (n, x_y_w_h_r)
            :param scores: 置信度分数数组，形状为 (n,)
            :param iou_thres: IoU阈值
            :return: 保留的边界框索引
        """
        # 假设 prediction 的形状为 (batch_size, num_detections, num_classes + 4 + num_masks)
        # masks 的形状为 (batch_size, num_detections, H, W)
        if iou_thres is None:
            iou_thres = self.iou_threshold
        scores = np.array(scores)
        obbox = np.array(obbox)
        if obbox.size > 0:
            # 按照置信度分数降序排序
            indices = np.argsort(scores)[::-1]
            obbox = obbox[indices]
            iou_score = self.probiou(obbox, obbox)
            iou_score = np.triu(iou_score, k=1)
            max_iou = np.max(iou_score, axis=0)
            pick = np.flatnonzero(max_iou < iou_thres)
            return indices[pick], max_iou[pick]
        return np.array([]), np.array([])

    def preprocess(self, img: np.ndarray):
        """
        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=(114, 114, 114))

        # 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_process = img[None] if len(img.shape) == 3 else img
        return img_process, ratio, (pad_w, pad_h)

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

        Args:
            preds (Numpy.ndarray): predictions come from ort.session.run().
            im0 (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 obb (xywh, theta, score, cls).
        """
        x = preds[0]  # Two outputs: predictions and protos

        # Transpose dim 1: (Batch_size, xywh_conf_cls_theta, Num_anchors) -> (Batch_size, Num_anchors, xywh_conf_cls_theta)
        x = np.einsum("bcn->bnc", x)

        # Predictions filtering by conf-threshold
        x = x[np.amax(x[..., 4:-1], axis=-1) > self.conf_threshold]
        x[..., 0:2] -= [pad_w, pad_h]
        x[..., 0:4] = x[..., 0:4] / ratio

        # Create a new matrix which merge these(xywh, theta, score, cls) into one
        x = np.c_[x[..., :4], x[..., -1:], np.amax(x[..., 4:-1], axis=-1), np.argmax(x[..., 4:-1], axis=-1)]

        boxes = np.copy(x[:, :5])  # [x, y, w, h, theta]
        c = x[:, 6:7] * max(im0.shape[0:2])  # 处理类别
        boxes[:, :2] += c
        # time1 = time.time()
        index_id, _ = self.nms_with_iou(boxes, x[:, 5])
        # print(f"{time.time() - time1}")
        if index_id.size > 0:
            x2 = x[index_id, ::]
            boxes = np.copy(x2[:, :5])  # [x, y, w, h, theta]
            index_id, iou_s = self.nms_with_iou(boxes, x2[:, 5], 0.8, re_oiu)
            x2 = x2[index_id, ::]
            if re_oiu:
                if iou_s.shape[0] > 23:
                    return x2, np.sum(iou_s) / iou_s.shape[0]
                return x2, np.sum(iou_s) / 23
            else:
                return x2, 0
        return None, 0

    # def find_center_point(self):

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

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

        Returns:
            None
        """
        # Draw rectangles and polygons
        im_canvas = im.copy()
        box_data = self.xywhr2xyxyxyxy(bboxes[:, :5])
        count = 0
        for box, conf, cls_ in zip(box_data, bboxes[:, 5], 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,
                0.7,
                (255, 0, 0),
                2,
                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)

    def xywhr2xyxyxyxy(self, 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)

    def get_score(self, obboxes: np.ndarray, iou_score, max_l: float = 223.6):
        """
        :param obboxes: 监测数据
        :param iou_score:
        :param max_l:
        :param img_size:
        :return:
        """
        """
        计算每个类别染色体最大长度的平均值， X, Y染色体不纳入计算
        参数格式: [[x, y, w, h, conf, 类别], ...]
        """
        if obboxes.shape[0] <= 10:
            return 0, obboxes.shape[0], np.array([0, 0]).tolist()

        # 定义所有点的坐标
        c_points = obboxes[:, 0:2]
        kd_tree = KDTree(c_points)
        dist_re = 0
        w_s = obboxes[:, 2:4].min(axis=-1)
        min_k = np.partition(w_s, 4)
        target_w = min_k[:5].mean()
        bending_s = []
        count = obboxes.shape[0]
        # if 45 <= count <= 47:
        #     count_score = 1.0
        # else:
        count_score = 1 - abs(count - 46) / 46.0  # 染色体数量
        len_score = 1 - abs(obboxes[:, 2:4].max() - max_l) / max_l  # 相对长度
        dist_s = []
        for item in obboxes:
            # 计算离散程度
            dist, index = kd_tree.query(item[0:2], k=2, p=1)
            dist_s.append(dist[1])

            # 取长边作为染色体长度(w<=h)
            if item[2] < item[3]:
                w, h = item[2], item[3]
            else:
                w, h = item[3], item[2]
            if w <= target_w:
                bending_s.append(0)
            else:
                bending_s.append((w - target_w) / h)  # < 1
        # dist_s = sorted(dist_s)

        bending_score = 1 - np.mean(bending_s)  # 弯曲度
        # dist_score = np.mean(dist_s[0:len(dist_s)//2]) / 5000.0  # 离散度
        # dist_score = np.mean(dist_s) / 500.0
        score = bending_score + len_score * 10 + count_score * 1000 + (1 - iou_score) * 100 + np.mean(
            obboxes[:, -2]) * 100  # + dist_score * 100
        d_c_point = np.mean(c_points, axis=0)  # - np.array([self.image_size[1], self.image_size[0]]) // 2
        # max_yx = np.max(c_points, axis=0)
        # min_yx = np.min(c_points, axis=0)
        # if d_c_point[0] < 0:
        #     if d_c_point[0] + min_yx[0] < 0:
        #         d_c_point[0] = (d_c_point[0] - min_yx[0]) // 2
        # else:
        #     if max_yx[0] - d_c_point[0] - self.image_size[1]  < 0:
        #         d_c_point[0] = (d_c_point[0] + max_yx[0] - self.image_size[1]) // 2
        #
        # if d_c_point[1] < 0:
        #     if d_c_point[1] + min_yx[1] < 0:
        #         d_c_point[1] = (d_c_point[1] - min_yx[1]) // 2
        # else:
        #     if max_yx[1] - d_c_point[1] - self.image_size[0]  < 0:
        #         d_c_point[1] = (d_c_point[1] + max_yx[1] - self.image_size[0]) // 2

        # 计算所有类别最大长度的平均值
        return int(score), int(count), d_c_point.astype(np.int32).tolist()  # 弯曲度，相对长度, 染色体数量,  离散度


