import cv2
import numpy as np
from ultralytics import YOLO
from . import config

class PersonDetector:
    def __init__(self):
        self.model = None
        self.load_model()
        
    def load_model(self):
        """加载YOLO模型"""
        try:
            self.model = YOLO(config.MODEL_PATH)
            print(f"YOLO模型加载成功: {config.MODEL_PATH}")
        except Exception as e:
            print(f"模型加载失败: {e}")
            print("正在下载YOLOv8n模型...")
            self.model = YOLO('yolov8n.pt')  # 会自动下载
            
    def detect_persons(self, frame):
        """
        检测图像中的人体
        返回: 检测结果列表，每个结果包含 [x1, y1, x2, y2, confidence]
        """
        if self.model is None:
            return []
        
        try:
            # 使用YOLO进行检测
            results = self.model(frame, verbose=False)
            
            persons = []
            for result in results:
                boxes = result.boxes
                if boxes is not None:
                    for box in boxes:
                        # 获取类别、置信度和坐标
                        cls = int(box.cls[0])
                        conf = float(box.conf[0])
                        
                        # 只保留人体检测结果
                        if cls == config.PERSON_CLASS_ID and conf >= config.CONFIDENCE_THRESHOLD:
                            # 获取边界框坐标
                            x1, y1, x2, y2 = box.xyxy[0].cpu().numpy()
                            persons.append([int(x1), int(y1), int(x2), int(y2), conf])
                            
            return persons
            
        except Exception as e:
            print(f"检测过程中出错: {e}")
            return []
    
    def draw_detections(self, frame, detections, distances):
        """
        在图像上绘制检测结果和距离信息
        """
        for i, detection in enumerate(detections):
            x1, y1, x2, y2, conf = detection
            
            # 绘制边界框
            cv2.rectangle(frame, (x1, y1), (x2, y2), config.BOX_COLOR, 2)
            
            # 准备显示文本
            person_id = f"Person #{i+1}"
            distance_text = f"Distance: {distances[i]}" if i < len(distances) else "Distance: N/A"
            conf_text = f"Conf: {conf:.2f}"
            
            # 计算文本位置
            text_y = y1 - 35 if y1 - 35 > 20 else y1 + 20
            
            # 绘制人员ID文本背景和文字
            id_text_size = cv2.getTextSize(person_id, config.FONT, config.FONT_SCALE, config.FONT_THICKNESS)[0]
            cv2.rectangle(frame, (x1, text_y - id_text_size[1] - 5), 
                         (x1 + id_text_size[0] + 10, text_y + 5), (255, 0, 0), -1)
            cv2.putText(frame, person_id, (x1 + 5, text_y), 
                       config.FONT, config.FONT_SCALE, config.TEXT_COLOR, config.FONT_THICKNESS)
            
            # 绘制距离文本背景和文字
            distance_text_y = text_y + 25
            distance_text_size = cv2.getTextSize(distance_text, config.FONT, config.FONT_SCALE, config.FONT_THICKNESS)[0]
            cv2.rectangle(frame, (x1, distance_text_y - distance_text_size[1] - 5), 
                         (x1 + distance_text_size[0] + 10, distance_text_y + 5), config.TEXT_BG_COLOR, -1)
            cv2.putText(frame, distance_text, (x1 + 5, distance_text_y), 
                       config.FONT, config.FONT_SCALE, config.TEXT_COLOR, config.FONT_THICKNESS)
            
            # 绘制置信度文本（在框的右上角）
            conf_text_size = cv2.getTextSize(conf_text, config.FONT, config.FONT_SCALE - 0.2, config.FONT_THICKNESS)[0]
            cv2.rectangle(frame, (x2 - conf_text_size[0] - 10, y1), 
                         (x2, y1 + conf_text_size[1] + 10), config.TEXT_BG_COLOR, -1)
            cv2.putText(frame, conf_text, (x2 - conf_text_size[0] - 5, y1 + conf_text_size[1] + 5), 
                       config.FONT, config.FONT_SCALE - 0.2, config.TEXT_COLOR, config.FONT_THICKNESS)
            
        return frame
    
    def get_model_info(self):
        """获取模型信息"""
        if self.model:
            return f"YOLO Model: {config.MODEL_PATH}"
        return "Model not loaded" 