"""
从main.py提取的YOLO识别测试文件
使用与main.py相同的识别逻辑
"""
import cv2
from utils.get_image import GetImage
from ultralytics import YOLO
from config import config_manager
import os

# 检查模型文件是否存在
model_path = r"best.pt"
model0_path = r"best0.pt"

if not os.path.exists(model_path):
    print(f"❌ 模型文件不存在: {model_path}")
    exit(1)
if not os.path.exists(model0_path):
    print(f"❌ 模型文件不存在: {model0_path}")
    exit(1)

# 加载YOLO模型（与main.py保持一致）
try:
    model = YOLO(model_path).to('cuda')
    model0 = YOLO(model0_path).to('cuda')
    print(f"✅ 模型加载成功: {model_path}")
    print(f"✅ 模型加载成功: {model0_path}")
except Exception as e:
    print(f"❌ 模型加载失败: {e}")
    exit(1)


def yolo_shibie(im_PIL, detections, model):
    """
    YOLO识别函数（与main.py中的实现完全一致）
    :param im_PIL: PIL图像对象
    :param detections: 检测结果字典
    :param model: YOLO模型
    :return: 更新后的detections字典
    """
    results = model(im_PIL)  # 目标检测
    for result in results:
        for i in range(len(result.boxes.xyxy)):
            left, top, right, bottom = result.boxes.xyxy[i]
            scalar_tensor = result.boxes.cls[i]
            value = scalar_tensor.item()
            label = result.names[int(value)]
            if label == 'center' or label == 'next' or label == 'boss' or label == 'zhaozi':
                player_x = int(left + (right - left) / 2)
                player_y = int(top + (bottom - top) / 2) + 30
                RW = [player_x, player_y]
                detections[label] = RW
            elif label == 'daojv' or label == 'gw':
                player_x = int(left + (right - left) / 2)
                player_y = int(top + (bottom - top) / 2) + 30
                RW = [player_x, player_y]
                detections[label].append(RW)
            elif label == 'npc1' or label == 'npc2' or label == 'npc3' or label == 'npc4':
                player_x = int(left + (right - left) / 2)
                player_y = int(bottom) + 30
                RW = [player_x, player_y]
                detections[label] = RW
    return detections


def main():
    """主函数"""
    print("="*60)
    print("YOLO识别测试（main.py逻辑）")
    print("="*60)
    
    # 从配置加载采集卡设置
    active_group = config_manager.get_active_group()
    
    if active_group is None:
        print("⚠️ 没有活动的配置组，使用默认设置")
        print("提示: 可以运行 python gui_config.py 设置配置")
        cam_index = 0
        width = 1920
        height = 1080
        use_model = model  # 默认使用model
    else:
        print(f"📋 使用配置组: {active_group['name']}")
        cam_index = active_group['camera_index']
        width = active_group['camera_width']
        height = active_group['camera_height']
        use_model = model0  # 城镇中使用model0
        print(f"   使用模型: model0 (best0.pt) - 用于城镇识别")
    
    print(f"   采集卡索引: {cam_index}")
    print(f"   分辨率: {width}x{height}")
    print()
    
    # 初始化采集卡
    print("🔧 正在初始化采集卡...")
    get_image = GetImage(
        cam_index=cam_index,
        width=width,
        height=height
    )
    
    if get_image.cap is None:
        print("❌ 采集卡初始化失败")
        print("请检查：")
        print("1. 采集卡是否正确连接")
        print("2. 采集卡索引是否正确")
        print("3. 采集卡驱动是否安装")
        return
    
    print("✅ 采集卡初始化成功")
    print("\n快捷键:")
    print("  'q' 或 ESC - 退出")
    print("  'm' - 切换模型 (model/model0)")
    print("  'd' - 显示/隐藏检测信息")
    print()
    
    try:
        frame_count = 0
        show_detections = True  # 是否显示检测信息
        current_model = use_model  # 当前使用的模型
        current_model_name = "model0" if use_model == model0 else "model"
        
        while True:
            # 获取帧
            frame_data = get_image.get_frame()
            
            if frame_data is None:
                print("⚠️ 无法获取帧，跳过...")
                continue
            
            # frame_data 是 [im_opencv_rgb, im_PIL] 格式
            im_opencv_rgb, im_PIL = frame_data
            
            if im_PIL is None:
                print("⚠️ PIL图像为空，跳过...")
                continue
            
            # 初始化检测结果字典
            detections = {
                'center': None, 'next': None,
                'npc1': None, 'npc2': None, 'npc3': None, 'npc4': None,
                'boss': None, 'zhaozi': None,
                'daojv': [], 'gw': []
            }
            
            # 执行YOLO检测（使用main.py的逻辑）
            detections = yolo_shibie(im_PIL, detections, current_model)
            
            # 获取绘制好框的图像用于显示
            try:
                results = current_model(im_PIL)
                result = results[0]
                frame_with_boxes_rgb = result.plot()
                frame_with_boxes_bgr = cv2.cvtColor(frame_with_boxes_rgb, cv2.COLOR_RGB2BGR)
            except Exception as e:
                print(f"⚠️ 绘制检测框失败: {e}")
                frame_with_boxes_bgr = cv2.cvtColor(im_opencv_rgb, cv2.COLOR_RGB2BGR)
            
            # 在图像上显示检测信息
            if show_detections:
                # 显示模型名称
                cv2.putText(frame_with_boxes_bgr, f"Model: {current_model_name}", 
                           (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
                
                # 显示检测到的目标
                y_offset = 60
                detected_items = []
                for key, value in detections.items():
                    if value is not None and value != []:
                        if key in ['daojv', 'gw']:
                            detected_items.append(f"{key}: {len(value)}个")
                        else:
                            detected_items.append(f"{key}: {value}")
                
                if detected_items:
                    text = f"Detected: {', '.join(detected_items[:5])}"  # 最多显示5个
                    if len(detected_items) > 5:
                        text += f" ... (+{len(detected_items)-5})"
                    cv2.putText(frame_with_boxes_bgr, text, 
                               (10, y_offset), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 0), 2)
            
            # 显示图像
            cv2.imshow("YOLO Detection (main.py logic)", frame_with_boxes_bgr)
            
            # 检查按键
            key = cv2.waitKey(1) & 0xFF
            if key in [27, ord('q'), ord('Q')]:
                print("\n用户退出")
                break
            elif key == ord('m') or key == ord('M'):
                # 切换模型
                if current_model == model:
                    current_model = model0
                    current_model_name = "model0"
                else:
                    current_model = model
                    current_model_name = "model"
                print(f"切换模型: {current_model_name}")
            elif key == ord('d') or key == ord('D'):
                show_detections = not show_detections
                print(f"显示检测信息: {'开启' if show_detections else '关闭'}")
            
            frame_count += 1
            if frame_count % 30 == 0:  # 每30帧打印一次
                print(f"📊 已处理 {frame_count} 帧 (模型: {current_model_name})")
                # 打印有检测到的目标
                detected_items = {k: v for k, v in detections.items() if v is not None and v != []}
                if detected_items:
                    print(f"   检测到: {detected_items}")
    
    except KeyboardInterrupt:
        print("\n\n用户中断测试")
    except Exception as e:
        print(f"\n❌ 测试过程中发生错误: {e}")
        import traceback
        traceback.print_exc()
    finally:
        # 清理资源
        get_image.release()
        cv2.destroyAllWindows()
        print("🔚 测试结束")


if __name__ == "__main__":
    main()

