"""
yolox基于pytorch框架的黑盒水印处理验证流程
"""
import os
import cv2
from watermark_verify.inference.yolox_inference import YOLOXInference
from watermark_verify.process.general_process_define import BlackBoxWatermarkProcessDefine
from watermark_verify.tools import parse_qrcode_label_file
from watermark_verify.tools.evaluate_tool import calculate_ciou


class ModelWatermarkProcessor(BlackBoxWatermarkProcessDefine):
    def __init__(self, model_filename):
        super(ModelWatermarkProcessor, self).__init__(model_filename)

    def process(self) -> bool:
        """
        根据流程定义进行处理，并返回模型标签验证结果
        :return: 模型标签验证结果
        """
        # 获取权重文件，使用触发集进行模型推理, 将推理结果与触发集预先二维码保存位置进行比对，在误差范围内则进行下一步，否则返回False
        cls_image_mapping = parse_qrcode_label_file.parse_labels(self.qrcode_positions_file)
        accessed_cls = set()
        for cls, images in cls_image_mapping.items():
            for image in images:
                image_path = os.path.join(self.trigger_dir, image)
                detect_result = self.detect_secret_label(image_path, self.qrcode_positions_file, (640, 640))
                if detect_result:
                    accessed_cls.add(cls)
                    break
        if not accessed_cls == set(cls_image_mapping.keys()):  # 所有的分类都检测出模型水印，模型水印检测结果为True
            return False

        verify_result = self.verify_label()  # 模型标签检测通过，进行标签验证
        return verify_result

    def detect_secret_label(self, image_path, watermark_txt, input_shape) -> bool:
        """
        对模型使用触发集进行检查，判断是否存在黑盒模型水印，如果对嵌入水印的图片样本正确率高于阈值，证明模型存在黑盒水印
        :param image_path: 输入图像路径
        :param watermark_txt: 水印标签文件路径
        :param input_shape: 模型输入图像大小，tuple
        :return: 检测结果
        """
        img = cv2.imread(image_path)
        height, width, channels = img.shape
        x_center, y_center, w, h, cls = parse_qrcode_label_file.load_watermark_info(watermark_txt, image_path)
        # 计算绝对坐标
        x1 = (x_center - w / 2) * width
        y1 = (y_center - h / 2) * height
        x2 = (x_center + w / 2) * width
        y2 = (y_center + h / 2) * height
        watermark_box = [x1, y1, x2, y2, cls]
        if len(watermark_box) == 0:
            return False

        dets = YOLOXInference(self.model_filename,input_size=input_shape).predict(image_path)

        if dets is not None:
            detect_result = detect_watermark(dets, watermark_box)
            return detect_result
        else:
            return False


def detect_watermark(dets, watermark_box, threshold=0.5):
    if dets.size == 0:  # 检查是否为空
        return False
    for box, score, cls in zip(dets[:, :4], dets[:, 4], dets[:, 5]):
        wm_box_coords = watermark_box[:4]
        wm_cls = watermark_box[4]
        if cls == wm_cls:
            ciou = calculate_ciou(box, wm_box_coords)
            if ciou > threshold:
                return True
    return False