"""
快速测试YOLO检测效果

检查现有YOLO模型能否准确检测新数据中的香菇
"""

import cv2
import sys
from pathlib import Path
import logging

# 添加src路径
sys.path.insert(0, str(Path(__file__).parent))

try:
    from .mushroom_detector import MushroomDetector
except ImportError:  # pragma: no cover - fallback for direct script usage
    from mushroom_detector import MushroomDetector

logging.basicConfig(level=logging.INFO, format='%(levelname)s: %(message)s')


def test_yolo_on_images(image_dir: str, yolo_model: str, num_samples: int = 5):
    """
    测试YOLO在新数据上的检测效果

    参数:
        image_dir: 图像目录
        yolo_model: YOLO模型路径
        num_samples: 测试样本数量
    """
    detector = MushroomDetector(yolo_model)

    image_files = list(Path(image_dir).glob('*.jpg'))[:num_samples]

    print(f"\n{'='*60}")
    print(f"YOLO检测测试 - 共测试 {len(image_files)} 张图片")
    print(f"{'='*60}\n")

    total_detections = 0
    low_confidence_count = 0

    for idx, image_file in enumerate(image_files, 1):
        print(f"[{idx}/{len(image_files)}] {image_file.name}")

        # 读取图像
        image = cv2.imread(str(image_file))
        if image is None:
            print(f"  ❌ 无法读取图像")
            continue

        # YOLO检测
        detections = detector.detect(image)

        if len(detections) == 0:
            print(f"  ❌ 未检测到香菇")
        else:
            total_detections += len(detections)
            print(f"  ✓ 检测到 {len(detections)} 个香菇")

            for i, det in enumerate(detections):
                conf = det['confidence']
                if conf < 0.7:
                    low_confidence_count += 1
                    print(f"    香菇#{i+1}: 置信度={conf:.3f} ⚠️ 偏低")
                else:
                    print(f"    香菇#{i+1}: 置信度={conf:.3f} ✓")

        # 可视化第一张图片
        if idx == 1:
            vis_image = image.copy()
            for det in detections:
                box = det['box']
                conf = det['confidence']
                x1, y1, x2, y2 = map(int, box)

                # 绘制检测框
                color = (0, 255, 0) if conf >= 0.7 else (0, 165, 255)  # 绿色/橙色
                cv2.rectangle(vis_image, (x1, y1), (x2, y2), color, 3)

                # 标注置信度
                label = f"{conf:.2f}"
                cv2.putText(vis_image, label, (x1, y1-10),
                           cv2.FONT_HERSHEY_SIMPLEX, 1.0, color, 2)

            # 保存可视化结果
            output_path = Path('../outputs') / 'yolo_test_sample.jpg'
            output_path.parent.mkdir(exist_ok=True)
            cv2.imwrite(str(output_path), vis_image)
            print(f"\n  可视化结果已保存: {output_path}")

        print()

    # 总结
    print(f"{'='*60}")
    print("检测总结")
    print(f"{'='*60}")
    print(f"总检测数: {total_detections}")
    print(f"平均每张: {total_detections / len(image_files):.1f} 个")
    print(f"低置信度(<0.7): {low_confidence_count} 个")

    if total_detections == 0:
        print("\n❌ YOLO模型无法检测新数据，需要重新训练！")
        return False
    elif low_confidence_count > total_detections * 0.3:
        print(f"\n⚠️  {low_confidence_count}/{total_detections} 检测置信度偏低")
        print("建议：考虑重新训练YOLO或调整置信度阈值")
        return False
    else:
        print("\n✓ YOLO模型在新数据上表现良好，可以使用")
        return True


if __name__ == "__main__":
    test_yolo_on_images(
        image_dir='../data/raw_images',
        yolo_model='../models/yolo_mushroom.pt',
        num_samples=10  # 测试10张图片
    )
