import cv2
import numpy as np
from ultralytics import YOLO
import os

class LicensePlateYOLO:
    """
    车牌YOLO检测器类
    负责加载YOLO pose模型并进行车牌检测和角点提取
    """
    
    def __init__(self, model_path=None):
        """
        初始化YOLO检测器
        
        参数:
            model_path: 模型文件路径，如果为None则使用默认路径
        """
        self.model = None
        self.model_path = model_path or self._get_default_model_path()
        self.class_names = {0: '蓝牌', 1: '绿牌'}
        self.load_model()
    
    def _get_default_model_path(self):
        """获取默认模型路径"""
        current_dir = os.path.dirname(__file__)
        return os.path.join(current_dir, "yolo11s-pose42.pt")
    
    def load_model(self):
        """
        加载YOLO pose模型
        
        返回:
            bool: 加载是否成功
        """
        try:
            if os.path.exists(self.model_path):
                self.model = YOLO(self.model_path)
                print(f"YOLO模型加载成功: {self.model_path}")
                return True
            else:
                print(f"模型文件不存在: {self.model_path}")
                return False
        except Exception as e:
            print(f"YOLO模型加载失败: {e}")
            return False
    
    def detect_license_plates(self, image, conf_threshold=0.5):
        """
        检测图像中的车牌
        
        参数:
            image: 输入图像 (numpy数组)
            conf_threshold: 置信度阈值
        
        返回:
            list: 检测结果列表，每个元素包含:
                - box: 边界框坐标 [x1, y1, x2, y2]
                - keypoints: 四个角点坐标 [[x1,y1], [x2,y2], [x3,y3], [x4,y4]]
                - confidence: 置信度
                - class_id: 类别ID (0=蓝牌, 1=绿牌)
                - class_name: 类别名称
        """
        if self.model is None:
            print("模型未加载")
            return []
        
        try:
            # 进行推理
            results = self.model(image, conf=conf_threshold, verbose=False)
            detections = []
            
            for result in results:
                # 检查是否有检测结果
                if result.boxes is None or result.keypoints is None:
                    continue
                
                # 提取检测信息
                boxes = result.boxes.xyxy.cpu().numpy()  # 边界框
                keypoints = result.keypoints.xy.cpu().numpy()  # 关键点
                confidences = result.boxes.conf.cpu().numpy()  # 置信度
                classes = result.boxes.cls.cpu().numpy()  # 类别
                
                # 处理每个检测结果
                for i in range(len(boxes)):
                    # 检查关键点数量是否为4个
                    if len(keypoints[i]) == 4:
                        class_id = int(classes[i])
                        detection = {
                            'box': boxes[i],
                            'keypoints': keypoints[i],
                            'confidence': confidences[i],
                            'class_id': class_id,
                            'class_name': self.class_names.get(class_id, '未知')
                        }
                        detections.append(detection)
                    else:
                        # 关键点不足4个，记录但标记为不完整
                        class_id = int(classes[i])
                        detection = {
                            'box': boxes[i],
                            'keypoints': keypoints[i] if len(keypoints[i]) > 0 else [],
                            'confidence': confidences[i],
                            'class_id': class_id,
                            'class_name': self.class_names.get(class_id, '未知'),
                            'incomplete': True  # 标记为不完整
                        }
                        detections.append(detection)
            
            return detections
            
        except Exception as e:
            print(f"检测过程中出错: {e}")
            return []
    
    def draw_detections(self, image, detections):
        """
        在图像上绘制检测结果
        
        参数:
            image: 输入图像
            detections: 检测结果列表
        
        返回:
            numpy.ndarray: 绘制了检测结果的图像
        """
        draw_image = image.copy()
        
        for i, detection in enumerate(detections):
            box = detection['box']
            keypoints = detection['keypoints']
            class_name = detection['class_name']
            confidence = detection['confidence']
            incomplete = detection.get('incomplete', False)
            
            # 绘制边界框
            x1, y1, x2, y2 = map(int, box)
            
            # 根据车牌类型选择颜色
            if class_name == '绿牌':
                box_color = (0, 255, 0)  # 绿色
            elif class_name == '蓝牌':
                box_color = (255, 0, 0)  # 蓝色
            else:
                box_color = (128, 128, 128)  # 灰色
            
            cv2.rectangle(draw_image, (x1, y1), (x2, y2), box_color, 2)
            
            # 绘制标签
            label = f"{class_name} {confidence:.2f}"
            if incomplete:
                label += " (不完整)"
            
            # 计算文本大小和位置
            font = cv2.FONT_HERSHEY_SIMPLEX
            font_scale = 0.6
            thickness = 2
            (text_width, text_height), _ = cv2.getTextSize(label, font, font_scale, thickness)
            
            # 绘制文本背景
            cv2.rectangle(draw_image, (x1, y1 - text_height - 10), 
                         (x1 + text_width, y1), box_color, -1)
            
            # 绘制文本
            cv2.putText(draw_image, label, (x1, y1 - 5), 
                       font, font_scale, (255, 255, 255), thickness)
            
            # 绘制关键点和连线
            if len(keypoints) >= 4 and not incomplete:
                # 四个角点完整，用黄色连线
                points = [(int(kp[0]), int(kp[1])) for kp in keypoints[:4]]
                
                # 绘制关键点
                for point in points:
                    cv2.circle(draw_image, point, 5, (0, 255, 255), -1)
                
                # 连接关键点形成四边形（按顺序连接）
                # 假设关键点顺序为: right_bottom, left_bottom, left_top, right_top
                for j in range(4):
                    cv2.line(draw_image, points[j], points[(j+1)%4], (0, 255, 255), 2)
                
            elif len(keypoints) > 0:
                # 关键点不完整，用红色标记现有点
                for kp in keypoints:
                    point = (int(kp[0]), int(kp[1]))
                    cv2.circle(draw_image, point, 5, (0, 0, 255), -1)
        
        return draw_image
    
    def correct_license_plate(self, image, keypoints, target_size=(240, 80)):
        """
        使用四个角点对车牌进行透视变换矫正
        
        参数:
            image: 原始图像
            keypoints: 四个角点坐标
            target_size: 目标尺寸 (width, height)
        
        返回:
            numpy.ndarray: 矫正后的车牌图像，如果失败返回None
        """
        if len(keypoints) != 4:
            return None
        
        try:
            # 将关键点转换为numpy数组
            src_points = np.array(keypoints, dtype=np.float32)
            
            # 定义目标矩形的四个角点
            # 假设关键点顺序为: right_bottom, left_bottom, left_top, right_top
            # 重新排序为标准顺序: left_top, right_top, right_bottom, left_bottom
            width, height = target_size
            dst_points = np.array([
                [0, 0],           # left_top
                [width, 0],       # right_top
                [width, height],  # right_bottom
                [0, height]       # left_bottom
            ], dtype=np.float32)
            
            # 重新排序源点以匹配目标点
            # 原顺序: right_bottom, left_bottom, left_top, right_top
            # 目标顺序: left_top, right_top, right_bottom, left_bottom
            reordered_src = np.array([
                src_points[2],  # left_top
                src_points[3],  # right_top
                src_points[0],  # right_bottom
                src_points[1]   # left_bottom
            ], dtype=np.float32)
            
            # 计算透视变换矩阵
            matrix = cv2.getPerspectiveTransform(reordered_src, dst_points)
            
            # 应用透视变换
            corrected = cv2.warpPerspective(image, matrix, target_size)
            
            return corrected
            
        except Exception as e:
            print(f"车牌矫正失败: {e}")
            return None
    
    def get_model_info(self):
        """
        获取模型信息
        
        返回:
            dict: 模型信息字典
        """
        if self.model is None:
            return {"status": "未加载", "path": self.model_path}
        
        return {
            "status": "已加载",
            "path": self.model_path,
            "model_type": "YOLO11 Pose",
            "classes": self.class_names
        }

def initialize_yolo_detector(model_path=None):
    """
    初始化YOLO检测器的便捷函数
    
    参数:
        model_path: 模型文件路径
    
    返回:
        LicensePlateYOLO: 初始化后的检测器实例
    """
    detector = LicensePlateYOLO(model_path)
    return detector

if __name__ == "__main__":
    # 测试代码
    detector = initialize_yolo_detector()
    print("检测器信息:", detector.get_model_info())