#!/usr/bin/env python3
"""
YOLO模型训练脚本
用于红外图像目标检测模型训练
"""

import argparse
import sys
from pathlib import Path

from ultralytics import YOLO


def train_yolo(
    data_yaml: str,
    model_size: str = "n",
    epochs: int = 100,
    imgsz: int = 640,
    batch: int = 16,
    device: str = "0",
    output_dir: str = "models/train",
    resume: bool = False,
) -> str:
    """
    训练YOLO模型

    Args:
        data_yaml: 数据集配置文件路径
        model_size: 模型大小 ('n', 's', 'm', 'l', 'x')
        epochs: 训练轮数
        imgsz: 输入图像大小
        batch: 批量大小
        device: 训练设备 ('0' for GPU, 'cpu' for CPU)
        output_dir: 输出目录
        resume: 是否从上次中断处继续训练

    Returns:
        最佳模型路径
    """
    print("=" * 50)
    print("YOLO 红外故障检测模型训练")
    print("=" * 50)
    print(f"数据集配置: {data_yaml}")
    print(f"模型大小: {model_size}")
    print(f"训练轮数: {epochs}")
    print(f"图像尺寸: {imgsz}")
    print(f"批量大小: {batch}")
    print(f"训练设备: {device}")
    print("=" * 50)

    # 加载预训练模型
    model = YOLO(f"yolov8{model_size}.pt")

    # 开始训练
    results = model.train(
        data=data_yaml,
        epochs=epochs,
        imgsz=imgsz,
        batch=batch,
        device=device,
        project=output_dir,
        name="infrared_fault",
        exist_ok=True,
        resume=resume,
        patience=20,
        save=True,
        save_period=10,
        plots=True,
        verbose=True,
    )

    best_model_path = f"{output_dir}/infrared_fault/weights/best.pt"
    last_model_path = f"{output_dir}/infrared_fault/weights/last.pt"

    print("\n" + "=" * 50)
    print("训练完成!")
    print(f"最佳模型: {best_model_path}")
    print(f"最终模型: {last_model_path}")
    print("=" * 50)

    return best_model_path


def validate_model(model_path: str, data_yaml: str, device: str = "0") -> dict:
    """
    验证模型性能

    Args:
        model_path: 模型路径
        data_yaml: 数据集配置路径
        device: 验证设备

    Returns:
        验证指标字典
    """
    print("=" * 50)
    print("模型验证")
    print("=" * 50)

    model = YOLO(model_path)
    metrics = model.val(data=data_yaml, device=device)

    results = {
        "mAP50": float(metrics.box.map50),
        "mAP50-95": float(metrics.box.map),
        "precision": float(metrics.box.mp),
        "recall": float(metrics.box.mr),
    }

    print("\n验证结果:")
    print(f"  mAP@0.5: {results['mAP50']:.4f}")
    print(f"  mAP@0.5:0.95: {results['mAP50-95']:.4f}")
    print(f"  Precision: {results['precision']:.4f}")
    print(f"  Recall: {results['recall']:.4f}")
    print("=" * 50)

    return results


def export_model(model_path: str, format: str = "onnx") -> str:
    """
    导出模型

    Args:
        model_path: 模型路径
        format: 导出格式 ('onnx', 'torchscript', 'openvino', etc.)

    Returns:
        导出后模型路径
    """
    print("=" * 50)
    print(f"导出模型: {model_path} -> {format}")
    print("=" * 50)

    model = YOLO(model_path)
    export_path = model.export(format=format)

    print(f"导出完成: {export_path}")
    return export_path


def main():
    parser = argparse.ArgumentParser(description="YOLO红外故障检测模型训练")
    subparsers = parser.add_subparsers(dest="command", help="命令")

    # Train command
    train_parser = subparsers.add_parser("train", help="训练模型")
    train_parser.add_argument("--data", type=str, required=True, help="数据集配置文件")
    train_parser.add_argument("--model-size", type=str, default="n",
                              choices=["n", "s", "m", "l", "x"],
                              help="模型大小 (default: n)")
    train_parser.add_argument("--epochs", type=int, default=100, help="训练轮数")
    train_parser.add_argument("--imgsz", type=int, default=640, help="图像尺寸")
    train_parser.add_argument("--batch", type=int, default=16, help="批量大小")
    train_parser.add_argument("--device", type=str, default="0", help="训练设备")
    train_parser.add_argument("--output", type=str, default="models/train", help="输出目录")
    train_parser.add_argument("--resume", action="store_true", help="继续训练")

    # Validate command
    val_parser = subparsers.add_parser("validate", help="验证模型")
    val_parser.add_argument("--model", type=str, required=True, help="模型路径")
    val_parser.add_argument("--data", type=str, required=True, help="数据集配置")
    val_parser.add_argument("--device", type=str, default="0", help="验证设备")

    # Export command
    export_parser = subparsers.add_parser("export", help="导出模型")
    export_parser.add_argument("--model", type=str, required=True, help="模型路径")
    export_parser.add_argument("--format", type=str, default="onnx",
                              choices=["onnx", "torchscript", "openvino", "tensorrt"],
                              help="导出格式")

    args = parser.parse_args()

    if args.command == "train":
        train_yolo(
            data_yaml=args.data,
            model_size=args.model_size,
            epochs=args.epochs,
            imgsz=args.imgsz,
            batch=args.batch,
            device=args.device,
            output_dir=args.output,
            resume=args.resume,
        )
    elif args.command == "validate":
        validate_model(args.model, args.data, args.device)
    elif args.command == "export":
        export_model(args.model, args.format)
    else:
        parser.print_help()
        sys.exit(1)


if __name__ == "__main__":
    main()