import os
import numpy as np
import cv2
from typing import List, Dict, Any
import torch

class Model:
    """
    YOLOv8 模型包装类 - 使用 Ultralytics YOLO
    """
    
    def __init__(self):
        """初始化YOLOv8模型"""
        # 获取当前文件所在目录路径
        model_dir = os.path.dirname(os.path.abspath(__file__))
        # 模型文件路径
        model_path = os.path.join(model_dir, "best.pt")
        
        print(f"正在加载YOLOv8模型: {model_path}")
        
        # 检查设备
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        print(f"使用设备: {self.device}")
        
        # 使用 Ultralytics YOLO 加载模型
        try:
            from ultralytics import YOLO
            self.model = YOLO(model_path)
            print("使用 Ultralytics YOLO 加载模型成功")
        except ImportError:
            raise ImportError("请安装 ultralytics: pip install ultralytics>=8.0.0")
        except Exception as e:
            raise Exception(f"加载YOLOv8模型失败: {str(e)}")
                
        # 加载类别名称
        self.classes = []
        classes_path = os.path.join(model_dir, "classes.txt")
        if os.path.exists(classes_path):
            with open(classes_path, 'r', encoding='utf-8') as f:
                self.classes = [line.strip() for line in f.readlines() if line.strip()]
            print(f"已加载 {len(self.classes)} 个类别")
        else:
            # 使用模型自带的类别信息
            if hasattr(self.model, 'names') and self.model.names:
                self.classes = list(self.model.names.values()) if isinstance(self.model.names, dict) else self.model.names
                print(f"使用模型自带类别，共 {len(self.classes)} 个类别")
            else:
                print("未找到类别文件，将使用数字索引作为类别名")
        
        # 设置识别参数
        self.conf_threshold = 0.25  # 置信度阈值
        self.img_size = 640        # 默认输入图像大小
        
        print("YOLOv8模型加载完成")
    
    def preprocess(self, image: np.ndarray) -> np.ndarray:
        """预处理图像 - YOLOv8会自动处理，这里直接返回"""
        return image
    
    def predict(self, image: np.ndarray) -> List[Dict[str, Any]]:
        """模型推理"""
        original_height, original_width = image.shape[:2]
        
        try:
            # YOLOv8推理
            results = self.model(
                image, 
                conf=self.conf_threshold,
                device=self.device,
                verbose=False
            )
            
            detections = []
            
            # 解析结果
            for result in results:
                # 获取检测框
                boxes = result.boxes
                
                if boxes is None or len(boxes) == 0:
                    continue
                
                # 遍历每个检测框
                for box in boxes:
                    # 获取坐标 (xyxy格式)
                    xyxy = box.xyxy[0].cpu().numpy()
                    x1, y1, x2, y2 = xyxy
                    
                    # 转换为归一化坐标 (x, y, w, h)
                    x = x1 / original_width
                    y = y1 / original_height
                    w = (x2 - x1) / original_width
                    h = (y2 - y1) / original_height
                    
                    # 获取置信度
                    conf = float(box.conf[0].cpu().numpy())
                    
                    # 获取类别ID
                    cls_id = int(box.cls[0].cpu().numpy())
                    
                    # 获取类别名称
                    class_name = f"cls{cls_id}"
                    if 0 <= cls_id < len(self.classes):
                        class_name = self.classes[cls_id]
                    
                    # 添加检测结果
                    if conf >= self.conf_threshold:
                        detections.append({
                            'bbox': (x, y, w, h),
                            'class_id': cls_id,
                            'confidence': conf
                        })
                        
            return detections
                
        except Exception as e:
            print(f"推理过程中出错: {str(e)}")
            import traceback
            traceback.print_exc()
            return []
    
    @property
    def applies_nms(self) -> bool:
        """模型是否内部应用了 NMS"""
        # YOLOv8会自动应用 NMS
        return True
    
    def close(self):
        """释放资源"""
        if hasattr(self, 'model'):
            # 删除模型以释放 GPU 内存
            del self.model
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
        print("YOLOv8模型已关闭")