@echo off
setlocal EnableExtensions EnableDelayedExpansion

set "ROOT=%~dp0.."
for %%I in ("%ROOT%") do set "ROOT=%%~fI"
set "EXPORTER=%ROOT%\build\msbuild\bin\x64\Release\TransformerTrainingExport.exe"
set "TRAINER=%ROOT%\tools\train_transformer.py"
set "PYTHONUTF8=1"
set "PYTHON_EXE="

if not exist "%EXPORTER%" (
    echo TransformerTrainingExport.exe is missing.
    echo Build projects\TransformerTrainingExport\TransformerTrainingExport.vcxproj as Release x64 first.
    exit /b 2
)

if not exist "%TRAINER%" (
    echo Trainer script is missing: "%TRAINER%"
    exit /b 2
)

if defined SZILASSI_PYTHON call :try_cuda_python "%SZILASSI_PYTHON%"
call :try_cuda_python "%ROOT%\.venv\Scripts\python.exe"
for /f "delims=" %%P in ('dir /b /ad /o-n "%LOCALAPPDATA%\Programs\Python\Python3*" 2^>nul') do call :try_cuda_python "%LOCALAPPDATA%\Programs\Python\%%P\python.exe"

if defined PYTHON_EXE (
    pushd "%ROOT%"
    "%PYTHON_EXE%" "%TRAINER%" --repository "%ROOT%" --exporter "%EXPORTER%" %*
    set "RESULT=!ERRORLEVEL!"
    popd
    exit /b !RESULT!
)

py -3 -c "import torch; raise SystemExit(0 if torch.cuda.is_available() else 1)" >nul 2>&1
if not errorlevel 1 (
    pushd "%ROOT%"
    py -3 "%TRAINER%" --repository "%ROOT%" --exporter "%EXPORTER%" %*
    set "RESULT=!ERRORLEVEL!"
    popd
    exit /b !RESULT!
)

python -c "import torch; raise SystemExit(0 if torch.cuda.is_available() else 1)" >nul 2>&1
if not errorlevel 1 (
    pushd "%ROOT%"
    python "%TRAINER%" --repository "%ROOT%" --exporter "%EXPORTER%" %*
    set "RESULT=!ERRORLEVEL!"
    popd
    exit /b !RESULT!
)

python3 -c "import torch; raise SystemExit(0 if torch.cuda.is_available() else 1)" >nul 2>&1
if not errorlevel 1 (
    pushd "%ROOT%"
    python3 "%TRAINER%" --repository "%ROOT%" --exporter "%EXPORTER%" %*
    set "RESULT=!ERRORLEVEL!"
    popd
    exit /b !RESULT!
)

echo No Python interpreter with CUDA-enabled PyTorch was found.
echo Install CUDA PyTorch or set SZILASSI_PYTHON to its python.exe path.
exit /b 2

:try_cuda_python
if defined PYTHON_EXE exit /b 0
if "%~1"=="" exit /b 0
if not exist "%~1" exit /b 0
"%~1" -c "import torch; raise SystemExit(0 if torch.cuda.is_available() else 1)" >nul 2>&1
if not errorlevel 1 set "PYTHON_EXE=%~1"
exit /b 0