"""
YOLOv8n 目标检测器训练脚本
基于准备好的 YOLO 数据集微调极轻量级模型，用于高精召回 672x480 背景图中的全部图标
"""

import logging
from pathlib import Path
import torch
from ultralytics import YOLO

logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
logger = logging.getLogger("train_detector")

DATASET_YAML = Path("tcaptcha/data/yolo/dataset.yaml")
MODEL_DIR = Path("tcaptcha/models")


def train(epochs: int = 40, batch: int = 8, imgsz: int = 640):
    if not DATASET_YAML.exists():
        logger.error("数据集配置文件未找到: %s，请先运行 prepare_yolo.py", DATASET_YAML)
        return

    MODEL_DIR.mkdir(parents=True, exist_ok=True)

    # 选择加速设备
    device = "mps" if torch.backends.mps.is_available() else "cpu"
    logger.info("采用训练加速设备: %s", device)

    # 加载预训练 YOLOv8n (极轻量，只有 3.2M 参数)
    model = YOLO("yolov8n.pt")

    logger.info("开始微调训练，轮数: %d, 批大小: %d, 分辨率: %d", epochs, batch, imgsz)

    results = model.train(
        data=str(DATASET_YAML.resolve()),
        epochs=epochs,
        batch=batch,
        imgsz=imgsz,
        device=device,
        project=str(MODEL_DIR / "runs"),
        name="detector_exp",
        exist_ok=True,
        workers=2,
        augment=True,
        degrees=45.0,  # 随机旋转增强抗噪能力
        flipud=0.2,
        fliplr=0.5,
        mosaic=1.0,
        verbose=True,
    )

    # 保存最佳权重到 models/yolo_detector.pt
    best_pt = MODEL_DIR / "runs" / "detector_exp" / "weights" / "best.pt"
    target_pt = MODEL_DIR / "yolo_detector.pt"
    if best_pt.exists():
        import shutil
        shutil.copy(best_pt, target_pt)
        logger.info("✓ 训练完成！最佳模型权重已保存至: %s", target_pt)
    else:
        logger.warning("未找到 best.pt 权重")


if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("--epochs", type=int, default=35, help="训练轮数 (默认: 35)")
    parser.add_argument("--batch", type=int, default=8, help="Batch size (默认: 8)")
    args = parser.parse_args()
    train(epochs=args.epochs, batch=args.batch)