@echo off
chcp 65001 >nul
cls

set "REPO_DIR=%~dp0"
pushd "%REPO_DIR%"

set "YOLO_LAUNCHER="
where yolo >nul 2>nul
if %ERRORLEVEL% equ 0 set "YOLO_LAUNCHER=yolo"

if not defined YOLO_LAUNCHER (
    py -3.11 -c "import ultralytics" >nul 2>nul
    if %ERRORLEVEL% equ 0 set "YOLO_LAUNCHER=py -3.11 -m ultralytics"
)

if not defined YOLO_LAUNCHER (
    python -c "import ultralytics" >nul 2>nul
    if %ERRORLEVEL% equ 0 set "YOLO_LAUNCHER=python -m ultralytics"
)

if not defined YOLO_LAUNCHER (
    echo [错误] 未找到可用的 Ultralytics 启动方式
    echo.
    echo 请先确保满足以下任一条件:
    echo   1. yolo 命令已加入 PATH
    echo   2. py -3.11 可运行并已安装 ultralytics
    echo   3. python 可运行并已安装 ultralytics
    popd
    pause
    exit /b 1
)

echo ============================================================
echo           训练鞋子检测模型 (YOLOv8 + 640x640)
echo ============================================================
echo.

:: 设置数据集路径
set "DATASET=%REPO_DIR%datasets\openimages-shoes-yolo\data.yaml"

:: 检查数据集是否存在
if not exist "%DATASET%" (
    echo [错误] 找不到数据集配置文件: %DATASET%
    echo.
    echo 请先下载数据集:
    echo   py -3.11 "%REPO_DIR%01_download_dataset.py" --source openimages --max-samples 5000
    popd
    pause
    exit /b 1
)

echo [信息] 数据集: %DATASET%
echo.

:: 选择模型
echo 选择模型:
echo   1. YOLOv8n (轻量级, 速度快)
echo   2. YOLOv8s (推荐, 速度和精度平衡)
echo   3. YOLOv8m (高精度, 较慢)
echo.
set /p MODEL_CHOICE="输入选择 (1-3, 默认 2): "

if "%MODEL_CHOICE%"=="" set MODEL_CHOICE=2
if "%MODEL_CHOICE%"=="1" (
    set MODEL=yolov8n.pt
    set DESC=YOLOv8n
)
if "%MODEL_CHOICE%"=="2" (
    set MODEL=yolov8s.pt
    set DESC=YOLOv8s (推荐)
)
if "%MODEL_CHOICE%"=="3" (
    set MODEL=yolov8m.pt
    set DESC=YOLOv8m
)

echo.
echo [信息] 使用模型: %DESC%
echo.

:: 训练参数
set EPOCHS=150
set IMGSZ=640
set BATCH=16

echo 训练参数:
echo   - Epochs: %EPOCHS%
echo   - Image Size: %IMGSZ%x%IMGSZ%
echo   - Batch Size: %BATCH%
echo   - Device: GPU (cuda:0)
echo.

echo ============================================================
echo                     开始训练
echo ============================================================
echo.

call %YOLO_LAUNCHER% detect train data="%DATASET%" model="%MODEL%" epochs=%EPOCHS% imgsz=%IMGSZ% batch=%BATCH% device=0

if %ERRORLEVEL% neq 0 (
    echo.
    echo [错误] 训练失败！
    popd
    pause
    exit /b 1
)

echo.
echo ============================================================
echo                     训练完成！
echo ============================================================
echo.
echo 模型保存在: runs/detect/train/weights/best.pt
echo.
echo 下一步: 运行 03_export_onnx.bat 导出 ONNX
echo.
popd
pause