#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
Ultralytics YOLO 测试脚本
- validate 模式：在 dataset.yaml 上评测 mAP、精度/召回
- predict  模式：对图片/文件夹/视频/摄像头进行推理并保存可视化
- test     模式：在指定 source 数据上评估模型指标（mAP、精度、召回等）
- annotation 模式：对图像进行检测并生成LabelMe格式的JSON标注文件（保存到与图像相同目录）

示例：
  # 在验证集上评测
  python3 detection/test_yolo.py \
    --weights runs/train/single_yolov8s_20250908_152257/weights/best.pt \
    --data ../datasets/yolo_dataset/dataset.yaml \
    --imgsz 640 \
    --mode validate

  # 对文件夹做推理并保存结果
  python3 detection/test_yolo.py \
    --weights runs/train/yolov8s_strong_20250909_114704/weights/best.pt \
    --source ./datasets/l0_9.9_sum/images/test \
    --imgsz 640 \
    --conf 0.25 \
    --mode predict \
    --project runs/predict \
    --name demo

  # 在指定数据上评估模型指标
  python3 detection/test_yolo.py \
    --weights runs/train/yolov8s_strong_20250909_114704/weights/best.pt \
    --source ./datasets/l0_9.9_sum/images/test \
    --imgsz 640 \
    --conf 0.25 \
    --mode test \
    --project runs/test \
    --name eval_test

  # 生成LabelMe格式的JSON标注文件，并像 predict 模式一样输出预测可视化
  python3 detection/test_yolo.py \
    --weights runs/train/yolov8s_strong_20250909_114704/weights/best.pt \
    --source ./7009-7509\
    --imgsz 640 \
    --conf 0.25 \
    --mode annotation \
    --project runs/annotations\
    --name annotations14
"""

import os
import sys
import argparse
import numpy as np
from pathlib import Path


def parse_args():
    ap = argparse.ArgumentParser(description="YOLO 模型测试/推理")
    ap.add_argument("--weights", required=True, help="权重路径，例如 runs/train/.../weights/best.pt")
    ap.add_argument("--mode", choices=["validate", "predict", "test", "annotation"], default="validate", help="测试模式")
    ap.add_argument("--data", default="", help="dataset.yaml 路径(用于 validate)")
    ap.add_argument("--source", default="", help="推理输入(用于 predict/test)：图片/文件夹/视频/rtsp/0(摄像头)")
    ap.add_argument("--imgsz", type=int, default=640, help="输入尺寸")
    ap.add_argument("--conf", type=float, default=0.25, help="置信度阈值(预测)")
    ap.add_argument("--iou", type=float, default=0.45, help="NMS IOU 阈值(预测)")
    ap.add_argument("--device", default="", help="CUDA 设备，如 '0' 或 '0,1'，留空自动")
    ap.add_argument("--project", default="runs/test", help="输出根目录")
    ap.add_argument("--name", default="", help="任务名(默认自动)")
    ap.add_argument("--save_txt", action="store_true", help="保存 YOLO txt 结果(预测)")
    ap.add_argument("--save_conf", action="store_true", help="在 txt 中保存置信度")
    ap.add_argument("--recursive", action="store_true", help="当 --source 为目录时递归查找(默认仅当前目录)")
    return ap.parse_args()


def list_images_in_dir(dir_path: Path, recursive: bool = False):
    exts = {".jpg", ".jpeg", ".png", ".bmp", ".tif", ".tiff"}
    if recursive:
        return [p for p in dir_path.rglob("*") if p.suffix.lower() in exts]
    else:
        return [p for p in dir_path.iterdir() if p.is_file() and p.suffix.lower() in exts]


def evaluate_on_source(model, source_path: str, args):
    """
    在指定的 source 数据上评估模型指标
    """
    from ultralytics.utils.metrics import ConfusionMatrix
    import torch
    import numpy as np
    
    src = Path(source_path)
    if not src.exists():
        print(f"[错误] 源路径不存在: {src}")
        return None
    
    # 收集所有图像文件
    if src.is_dir():
        image_files = list_images_in_dir(src, recursive=args.recursive)
        if not image_files:
            print(f"[错误] 目录下未找到图像: {src}")
            return None
    else:
        image_files = [src]
    
    print(f"[测试] 找到 {len(image_files)} 个图像文件")
    
    # 创建临时 dataset.yaml 用于评估
    import tempfile
    import yaml
    
    # 假设标签文件在对应的 labels 目录中
    labels_dir = src.parent.parent / "labels" / src.name
    if not labels_dir.exists():
        print(f"[警告] 未找到标签目录: {labels_dir}")
        print("[测试] 将仅进行推理，无法计算精确指标")
        return run_inference_only(model, image_files, args)
    
    # 创建临时 dataset.yaml
    temp_yaml = {
        'path': str(src.parent),
        'train': '',  # 不使用
        'val': str(src.relative_to(src.parent)),  # 使用 source 作为验证集
        'test': '',
        'nc': 1,  # 假设单类检测
        'names': ['fish']  # 类别名称
    }
    
    with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml', delete=False) as f:
        yaml.dump(temp_yaml, f)
        temp_yaml_path = f.name
    
    try:
        # 使用 model.val 进行评估
        print("[测试] 开始评估...")
        metrics = model.val(
            data=temp_yaml_path,
            imgsz=args.imgsz,
            conf=args.conf,
            iou=args.iou,
            device=args.device if args.device else None,
            project=args.project,
            name=args.name or "test_eval",
            verbose=True,
            save_json=True,  # 保存详细结果
        )
        
        # 打印详细指标
        print("\n" + "="*50)
        print("评估结果:")
        print("="*50)

        # print number of images
        print(f"总图像数:    {len(image_files)}")
        
        if hasattr(metrics, 'box'):
            box_metrics = metrics.box
            print(f"mAP@0.5:     {box_metrics.map50:.4f}")
            print(f"mAP@0.5:0.95: {box_metrics.map:.4f}")
            print(f"Precision:   {box_metrics.mp:.4f}")
            print(f"Recall:      {box_metrics.mr:.4f}")
            print(f"F1-Score:    {2 * box_metrics.mp * box_metrics.mr / (box_metrics.mp + box_metrics.mr):.4f}")
        
        if hasattr(metrics, 'speed'):
            speed = metrics.speed
            print(f"推理速度:    {speed['inference']:.2f} ms/image")
            print(f"NMS速度:     {speed['postprocess']:.2f} ms/image")
        
        print("="*50)
        
        return metrics
        
    except Exception as e:
        print(f"[错误] 评估失败: {e}")
        print("[测试] 回退到仅推理模式...")
        return run_inference_only(model, image_files, args)
    finally:
        # 清理临时文件
        try:
            os.unlink(temp_yaml_path)
        except:
            pass


def run_inference_only(model, image_files, args):
    """
    仅进行推理，计算基本统计信息
    """
    print("[测试] 仅推理模式 - 无法计算精确指标")
    
    total_detections = 0
    total_images = len(image_files)
    confidences = []
    
    for img_path in image_files:
        try:
            results = model.predict(
                source=str(img_path),
                imgsz=args.imgsz,
                conf=args.conf,
                iou=args.iou,
                device=args.device if args.device else None,
                verbose=False,
                save=True,  # 同步保存Ultralytics自带的可视化结果
                project=args.project,
                name=args.name or "annotations",
            )
            
            if results and len(results) > 0:
                res = results[0]
                if hasattr(res, 'boxes') and res.boxes is not None:
                    num_dets = len(res.boxes)
                    total_detections += num_dets
                    
                    # 收集置信度
                    for box in res.boxes:
                        conf = float(box.conf[0].tolist())
                        confidences.append(conf)
                        
        except Exception as e:
            print(f"[警告] 处理图像失败 {img_path}: {e}")
    
    # 打印统计信息
    print("\n" + "="*50)
    print("推理统计:")
    print("="*50)
    print(f"总图像数:    {total_images}")
    print(f"总检测数:    {total_detections}")
    print(f"平均检测数:  {total_detections/total_images:.2f}")
    
    if confidences:
        confidences = np.array(confidences)
        print(f"平均置信度:  {confidences.mean():.4f}")
        print(f"最高置信度:  {confidences.max():.4f}")
        print(f"最低置信度:  {confidences.min():.4f}")
    
    print("="*50)
    print("[注意] 这是仅推理模式，无法计算 mAP 等精确指标")
    print("       如需精确评估，请确保有对应的标签文件")
    
    return None


def generate_labelme_annotations(model, source_path: str, args):
    """
    对图像进行检测并生成LabelMe格式的JSON标注文件（保存到 runs/<project>/<name>/），
    可视化结果由 Ultralytics 的 predict 保存（与 predict 模式一致）。
    
    Args:
        model: YOLO模型
        source_path: 图像目录路径
        args: 命令行参数
    """
    import json
    import cv2
    import numpy as np
    
    src = Path(source_path)
    if not src.exists():
        print(f"[错误] 源路径不存在: {src}")
        return
    
    # 收集所有图像文件
    if src.is_dir():
        image_files = list_images_in_dir(src, recursive=args.recursive)
        if not image_files:
            print(f"[错误] 目录下未找到图像: {src}")
            return
    else:
        image_files = [src]
    
    print(f"[标注] 找到 {len(image_files)} 个图像文件")
    
    # 创建输出目录 runs/<project>/<name>/
    project_dir = Path(args.project) #/ (args.name or "annotations")
    project_dir.mkdir(parents=True, exist_ok=True)
    output_dir = Path(args.project) / (args.name or "annotations")
    
    # 流式推理并即时写出标注，避免一次性占用大量显存/内存
    print(f"[标注] 流式预测并保存可视化到: {output_dir}")
    processed_count = 0
    failed_count = 0

    try:
        for res in model.predict(
            source=[str(p) for p in image_files],
            imgsz=args.imgsz,
            conf=args.conf,
            iou=args.iou,
            device=args.device if args.device else None,
            verbose=False,
            save=True,
            project=args.project,
            name=args.name or "annotations",
            stream=True,
        ):
            try:
                # 结果对象自带路径与原图尺寸
                img_path = Path(getattr(res, 'path', ''))
                if not img_path:
                    # 回退：无法获得路径时跳过
                    failed_count += 1
                    continue

                if hasattr(res, 'orig_shape') and res.orig_shape is not None:
                    img_height, img_width = int(res.orig_shape[0]), int(res.orig_shape[1])
                else:
                    original_img = cv2.imread(str(img_path))
                    if original_img is None:
                        print(f"[错误] 无法读取图像: {img_path}")
                        failed_count += 1
                        continue
                    img_height, img_width = original_img.shape[:2]

                labelme_data = {
                    "version": "5.8.3",
                    "flags": {},
                    "shapes": [],
                    "imagePath": img_path.name,
                    "imageData": None,
                    "imageHeight": img_height,
                    "imageWidth": img_width
                }

                if hasattr(res, 'boxes') and res.boxes is not None:
                    for box in res.boxes:
                        x1, y1, x2, y2 = [float(v) for v in box.xyxy[0].tolist()]
                        conf = float(box.conf[0].tolist())
                        cls = int(box.cls[0].tolist())
                        shape = {
                            "label": "鱿鱼",
                            "points": [[x1, y1], [x2, y2]],
                            "group_id": None,
                            "description": f"confidence: {conf:.3f}",
                            "shape_type": "rectangle",
                            "flags": {},
                            "mask": None
                        }
                        labelme_data["shapes"].append(shape)

                json_filename = img_path.stem + ".json"
                json_path = output_dir / json_filename
                with open(json_path, 'w', encoding='utf-8') as f:
                    json.dump(labelme_data, f, ensure_ascii=False, indent=2)

                print(f"[成功] {img_path.name} -> {json_filename} (检测到 {len(labelme_data['shapes'])} 个目标)")
                processed_count += 1
            except Exception as e:
                print(f"[错误] 处理图像失败 {getattr(res, 'path', 'unknown')}: {e}")
                failed_count += 1
    except Exception as e:
        print(f"[错误] 批处理预测失败: {e}")
    
    # 输出统计信息
    print("\n" + "="*60)
    print("标注生成完成:")
    print("="*60)
    print(f"成功处理: {processed_count} 个图像")
    print(f"处理失败: {failed_count} 个图像")
    print(f"输出目录: {output_dir}")
    print("生成文件:")
    print("  - JSON标注文件: *.json")
    print("  - 预测可视化: 由 Ultralytics 自动保存 (与 predict 模式一致)")
    print("="*60)


def main():
    args = parse_args()
    try:
        from ultralytics import YOLO
    except Exception as e:
        print("[错误] 未找到 ultralytics，请先: pip install ultralytics")
        print(e)
        sys.exit(1)
    
    # 检查matplotlib兼容性
    try:
        import matplotlib
        import numpy as np
        # 测试matplotlib和numpy的兼容性
        matplotlib.use('Agg')  # 使用非交互式后端
        import matplotlib.pyplot as plt
        print(f"[信息] Matplotlib版本: {matplotlib.__version__}")
        print(f"[信息] NumPy版本: {np.__version__}")
    except Exception as e:
        print(f"[警告] Matplotlib/NumPy兼容性问题: {e}")
        print("[建议] 尝试升级matplotlib: pip install matplotlib --upgrade")
        # 继续执行，但可能会在后续遇到问题

    model = YOLO(args.weights)

    if args.mode == "validate":
        if not args.data:
            print("[错误] validate 模式需要 --data 指向 dataset.yaml")
            sys.exit(1)
        metrics = model.val(
            data=args.data,
            imgsz=args.imgsz,
            device=args.device if args.device else None,
            project=args.project,
            name=args.name or "val",
            iou=args.iou,
            verbose=True,
        )
        # metrics 包含各项指标，可按需打印
        print("评测完成。mAP50:", getattr(metrics, 'box', None).map50 if hasattr(metrics, 'box') else None)
    elif args.mode == "test":
        if not args.source:
            print("[错误] test 模式需要 --source 输入")
            sys.exit(1)
        metrics = evaluate_on_source(model, args.source, args)
        if metrics:
            print("测试完成。")
        else:
            print("测试完成（仅推理模式）。")
    elif args.mode == "annotation":
        if not args.source:
            print("[错误] annotation 模式需要 --source 输入")
            sys.exit(1)
        generate_labelme_annotations(model, args.source, args)
        print("标注生成完成。")
    else:  # predict mode
        if not args.source:
            print("[错误] predict 模式需要 --source 输入")
            sys.exit(1)

        src = Path(args.source)
        sources = None
        if src.is_dir():
            # 严格限制在该目录（默认非递归）
            sources = list_images_in_dir(src, recursive=args.recursive)
            if not sources:
                print(f"[错误] 目录下未找到图像: {src}")
                sys.exit(1)
        else:
            # 单个文件或其他流式输入，直接传递
            sources = [str(src)]

        last_save_dir = None
        # 使用流式推理避免一次性将所有结果保存在内存/显存中
        for res in model.predict(
            source=str(src) if src.is_dir() else [str(src)],
            imgsz=args.imgsz,
            conf=args.conf,
            iou=args.iou,
            device=args.device if args.device else None,
            project=args.project,
            name=args.name or "pred",
            save=True,
            save_txt=args.save_txt,
            save_conf=args.save_conf,
            verbose=True,
            stream=True,
        ):
            last_save_dir = getattr(res, 'save_dir', None)
        print("推理完成。结果目录:", last_save_dir or args.project)


if __name__ == "__main__":
    main()