@echo off
setlocal

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 [ERROR] No usable Ultralytics launcher was found
    echo.
    echo Please make sure one of the following works:
    echo   1. yolo in PATH
    echo   2. py -3.11 with ultralytics installed
    echo   3. python with ultralytics installed
    popd
    pause
    exit /b 1
)

echo ============================================================
echo        Stage 2 Fine-tuning on PPE shoe subset
echo ============================================================
echo.

set "DATASET=%REPO_DIR%datasets\ppe-shoes\data.yaml"
set "BASE_MODEL=%REPO_DIR%runs\detect\train\weights\best.pt"

if not exist "%DATASET%" (
    echo [ERROR] PPE shoe subset not found: %DATASET%
    echo Run:
    echo   python "%REPO_DIR%05_prepare_ppe_shoe_subset.py"
    popd
    pause
    exit /b 1
)

if not exist "%BASE_MODEL%" (
    echo [ERROR] Base model not found: %BASE_MODEL%
    echo Run 02_train.bat first to create the base checkpoint.
    popd
    pause
    exit /b 1
)

echo [INFO] Dataset: %DATASET%
echo [INFO] Base model: %BASE_MODEL%
echo.

set "EPOCHS=50"
set "IMGSZ=640"
set "BATCH=16"

echo Fine-tune params:
echo   Epochs: %EPOCHS%
echo   Image Size: %IMGSZ%x%IMGSZ%
echo   Batch Size: %BATCH%
echo   Device: GPU (cuda:0)
echo.

echo ============================================================
echo                    Start fine-tuning
echo ============================================================
echo.

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

if %ERRORLEVEL% neq 0 (
    echo.
    echo [ERROR] Fine-tuning failed.
    popd
    pause
    exit /b 1
)

echo.
echo ============================================================
echo                    Fine-tuning complete
echo ============================================================
echo.
popd
pause