@echo off
setlocal EnableExtensions EnableDelayedExpansion

set "BASE=%~dp0"
set "ROOT=%LOCALAPPDATA%\MCAT_UNET"
set "ENV_NAME=py311"
set "ENV_DIR=%ROOT%\%ENV_NAME%"
set "PY_ROOT=%ENV_DIR%\python.exe"
set "PY_VENV=%ENV_DIR%\Scripts\python.exe"
set "PY="
set "SCRIPT=%BASE%mcat_learn_unet1d_1_3.py"
set "LOG=%BASE%setup_unet1d13_env.log"
set "PY_INSTALL_LOG=%BASE%python_install_unet1d13.log"
set "DOWNLOAD_DIR=%BASE%_downloads"
set "HEADLESS_SCRIPT=%BASE%mcat_train_selected_unet1d_headless.py"
set "PYTHON_VERSION=3.11.9"
set "PYTHON_INSTALLER=python-%PYTHON_VERSION%-amd64.exe"
set "PYTHON_URL=https://www.python.org/ftp/python/%PYTHON_VERSION%/%PYTHON_INSTALLER%"
set "PYTHON_INSTALLER_PATH=%DOWNLOAD_DIR%\%PYTHON_INSTALLER%"

echo MCAT U-Net v1.3 environment setup
echo Base: %BASE%
echo Python env: %ENV_DIR%
echo Log : %LOG%
echo Python installer log: %PY_INSTALL_LOG%
echo.

if not exist "%SCRIPT%" (
  echo Script not found: %SCRIPT%
  pause
  exit /b 1
)

if not exist "%HEADLESS_SCRIPT%" (
  echo Headless script not found: %HEADLESS_SCRIPT%
  pause
  exit /b 1
)

if not exist "%ROOT%" mkdir "%ROOT%"
if not exist "%DOWNLOAD_DIR%" mkdir "%DOWNLOAD_DIR%"

set "GPU_NAMES="
for /f "usebackq delims=" %%G in (`powershell -NoProfile -ExecutionPolicy Bypass -Command "try { (Get-CimInstance Win32_VideoController -ErrorAction Stop).Name } catch { try { (Get-WmiObject Win32_VideoController -ErrorAction Stop).Name } catch { '' } }"`) do (
  if not "%%G"=="" set "GPU_NAMES=!GPU_NAMES! %%G"
)

if "%GPU_NAMES%"=="" (
  for /f "skip=1 delims=" %%G in ('wmic path win32_VideoController get name 2^>nul') do (
    if not "%%G"=="" set "GPU_NAMES=!GPU_NAMES! %%G"
  )
)

echo Detected GPU(s): !GPU_NAMES!
echo Detected GPU(s): !GPU_NAMES! > "%LOG%"

set "TARGET=CPU"
set "TORCH_INSTALL=CPU"
echo !GPU_NAMES! | findstr /I "NVIDIA GeForce RTX Quadro" >nul
if not errorlevel 1 (
  set "TARGET=CUDA"
  set "TORCH_INSTALL=CUDA128"
)

if "%TARGET%"=="CPU" (
  echo !GPU_NAMES! | findstr /I "AMD Radeon Intel Arc" >nul
  if not errorlevel 1 (
    set "TARGET=DIRECTML"
    set "TORCH_INSTALL=DIRECTML"
  )
)

if "%TARGET%"=="CUDA" (
  set "NVIDIA_DRIVER="
  set "NVIDIA_CUDA_REPORTED="
  for /f "tokens=1 delims=," %%D in ('nvidia-smi --query-gpu^=driver_version --format^=csv^,noheader 2^>nul') do (
    if not defined NVIDIA_DRIVER set "NVIDIA_DRIVER=%%D"
  )
  for /f "tokens=1 delims=," %%C in ('nvidia-smi --query-gpu^=cuda_version --format^=csv^,noheader 2^>nul') do (
    if not defined NVIDIA_CUDA_REPORTED set "NVIDIA_CUDA_REPORTED=%%C"
  )
  echo NVIDIA driver: !NVIDIA_DRIVER!
  echo NVIDIA reported CUDA: !NVIDIA_CUDA_REPORTED!
  echo NVIDIA driver: !NVIDIA_DRIVER!>> "%LOG%"
  echo NVIDIA reported CUDA: !NVIDIA_CUDA_REPORTED!>> "%LOG%"

  echo !GPU_NAMES! | findstr /I "RTX 50 RTX50 5060 5070 5080 5090" >nul
  if not errorlevel 1 (
    echo RTX 50 series detected. PyTorch CUDA 12.8 wheels will be installed.
    echo Please keep the NVIDIA display driver on R570 or newer for RTX 50 series.
  ) else (
    echo NVIDIA GPU detected. PyTorch CUDA 12.8 wheels will be installed.
  )
)

if "%TARGET%"=="DIRECTML" (
  echo AMD/Intel GPU path selected. torch-directml will be installed.
  echo DirectML does not install CUDA.
)

if "%TARGET%"=="CPU" (
  echo No NVIDIA CUDA GPU or DirectML candidate was detected. CPU PyTorch will be installed.
)

echo Selected target: %TARGET%
echo Selected target: %TARGET%>> "%LOG%"
echo.

call :SELECT_ENV_PY
if defined PY (
  "!PY!" -B -c "import tkinter" >nul 2>nul
  if errorlevel 1 (
    echo Existing Python does not have tkinter. Rebuilding: %ENV_DIR%
    rmdir /s /q "%ENV_DIR%"
    set "PY="
  ) else (
    echo Existing Python found: !PY!
    set /p REBUILD=Rebuild it from scratch? This removes %ENV_DIR% only. [y/N]: 
    if /I "!REBUILD!"=="Y" (
      rmdir /s /q "%ENV_DIR%"
      if exist "%ENV_DIR%" (
        echo Failed to remove existing Python folder.
        pause
        exit /b 1
      )
      set "PY="
    ) else (
      echo Keeping existing Python. Checking required ONNX export packages...
      "!PY!" -m pip install onnx onnxscript
      if errorlevel 1 goto INSTALL_FAILED
      goto VERIFY
    )
  )
)

if exist "%ENV_DIR%" (
  echo Removing incomplete Python folder: %ENV_DIR%
  rmdir /s /q "%ENV_DIR%"
)

call :FIND_BASE_PY
if defined BASE_PY (
  echo Existing Python 3.11 found: !BASE_PY!
  echo Creating isolated MCAT environment from existing Python...
  call :CREATE_VENV
  if errorlevel 1 (
    echo Failed to create Python virtual environment from: !BASE_PY!
    pause
    exit /b 1
  )
  goto PYTHON_READY
)

call :DETECT_STALE_PY311
if defined STALE_PY311 (
  echo Broken Python 3.11.9 Windows Installer registration was detected.
  echo Python 3.11 executable was not found, but Python 3.11.9 components remain registered.
  echo.
  echo Please remove Python 3.11.9 from Windows Settings first:
  echo   Settings ^> Apps ^> Installed apps
  echo   Remove all entries named "Python 3.11.9 ... (64-bit)".
  echo.
  echo If normal uninstall fails, use Microsoft's "Program Install and Uninstall troubleshooter"
  echo to remove the stale Python 3.11.9 entries, then run this setup again.
  echo.
  pause
  exit /b 1
)

echo Downloading Python %PYTHON_VERSION% installer...
if not exist "%PYTHON_INSTALLER_PATH%" (
  powershell -NoProfile -ExecutionPolicy Bypass -Command "[Net.ServicePointManager]::SecurityProtocol=[Net.SecurityProtocolType]::Tls12; Invoke-WebRequest -Uri '%PYTHON_URL%' -OutFile '%PYTHON_INSTALLER_PATH%'"
  if errorlevel 1 (
    echo Failed to download Python installer.
    pause
    exit /b 1
  )
) else (
  echo Using cached installer: %PYTHON_INSTALLER_PATH%
)

echo Installing Python with tkinter to: %ENV_DIR%
if exist "%PY_INSTALL_LOG%" del "%PY_INSTALL_LOG%"
start /wait "" "%PYTHON_INSTALLER_PATH%" /quiet InstallAllUsers=0 TargetDir="%ENV_DIR%" Include_pip=1 Include_tcltk=1 Include_launcher=0 AssociateFiles=0 Shortcuts=0 PrependPath=0 Include_test=0 SimpleInstall=0 /log "%PY_INSTALL_LOG%"
set "PY_EXIT=%ERRORLEVEL%"
echo Python installer exit code: %PY_EXIT%
if not "%PY_EXIT%"=="0" (
  echo Python installer failed. See:
  echo %PY_INSTALL_LOG%
  pause
  exit /b 1
)

call :SELECT_ENV_PY
if not defined PY (
  echo Python was not created: %PY_ROOT%
  echo Installer exit code was 0, but python.exe is missing.
  echo A stale Python 3.11 installer registration may exist.
  echo Trying to create the MCAT environment from the registered Python instead...
  call :FIND_BASE_PY
  if defined BASE_PY (
    echo Existing Python 3.11 found after installer check: !BASE_PY!
    call :CREATE_VENV
    if errorlevel 1 (
      echo Failed to create Python virtual environment from: !BASE_PY!
      pause
      exit /b 1
    )
    goto PYTHON_READY
  )
)

if not defined PY (
  echo Python was not created: %PY_ROOT%
  echo Check whether Windows Installer blocked this environment path.
  if exist "%PY_INSTALL_LOG%" (
    echo Last install log lines:
    powershell -NoProfile -ExecutionPolicy Bypass -Command "Get-Content -LiteralPath '%PY_INSTALL_LOG%' -Tail 40"
  )
  pause
  exit /b 1
)

:PYTHON_READY
call :SELECT_ENV_PY
if not defined PY (
  echo Python environment was not created.
  pause
  exit /b 1
)

echo Checking tkinter...
"!PY!" -B -c "import tkinter; print('tkinter ok')"
if errorlevel 1 (
  echo tkinter check failed. Python was installed but Tcl/Tk is missing.
  if exist "%PY_INSTALL_LOG%" powershell -NoProfile -ExecutionPolicy Bypass -Command "Get-Content -LiteralPath '%PY_INSTALL_LOG%' -Tail 40"
  pause
  exit /b 1
)

echo Upgrading pip tools...
"!PY!" -m pip install --upgrade pip setuptools wheel
if errorlevel 1 (
  echo Failed to upgrade pip tools.
  pause
  exit /b 1
)

if "%TORCH_INSTALL%"=="CUDA128" (
  echo Installing PyTorch CUDA 12.8 build...
  "!PY!" -m pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu128
  if errorlevel 1 goto INSTALL_FAILED
  "!PY!" -m pip install numpy matplotlib onnx onnxscript
  if errorlevel 1 goto INSTALL_FAILED
  goto VERIFY
)

if "%TORCH_INSTALL%"=="DIRECTML" (
  echo Installing torch-directml path...
  "!PY!" -m pip install "numpy<2" matplotlib onnx onnxscript torch-directml
  if errorlevel 1 goto INSTALL_FAILED
  goto VERIFY
)

echo Installing CPU PyTorch...
"!PY!" -m pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu
if errorlevel 1 goto INSTALL_FAILED
"!PY!" -m pip install numpy matplotlib onnx onnxscript
if errorlevel 1 goto INSTALL_FAILED
goto VERIFY

:INSTALL_FAILED
echo Package installation failed. Check your internet connection and the log above.
pause
exit /b 1

:VERIFY
echo.
echo Verifying environment...
set "KMP_DUPLICATE_LIB_OK=TRUE"
"!PY!" -B -c "import sys, importlib.util, tkinter, torch; print('python', sys.version.split()[0]); print('tkinter ok'); print('torch', torch.__version__); print('cuda_build', torch.version.cuda); print('cuda_available', torch.cuda.is_available()); print('cuda_device_count', torch.cuda.device_count()); print('directml_installed', importlib.util.find_spec('torch_directml') is not None); import numpy, matplotlib, onnx, onnxscript; print('numpy', numpy.__version__); print('matplotlib', matplotlib.__version__); print('onnx', onnx.__version__); print('onnxscript ok')"
if errorlevel 1 (
  echo Verification failed.
  pause
  exit /b 1
)

"!PY!" -B -m py_compile "%SCRIPT%"
if errorlevel 1 (
  echo mcat_learn_unet1d_1_3.py compile check failed.
  pause
  exit /b 1
)

"!PY!" -B -m py_compile "%HEADLESS_SCRIPT%"
if errorlevel 1 (
  echo mcat_train_selected_unet1d_headless.py compile check failed.
  pause
  exit /b 1
)

echo.
echo Setup finished.
echo Run Learn_unet1d_1_3_AutoGPU.bat to start the trainer.
pause
endlocal
exit /b 0

:SELECT_ENV_PY
set "PY="
if exist "%PY_ROOT%" (
  set "PY=%PY_ROOT%"
  exit /b 0
)
if exist "%PY_VENV%" (
  set "PY=%PY_VENV%"
  exit /b 0
)
exit /b 0

:FIND_BASE_PY
set "BASE_PY="
set "CANDIDATE="
for /f "delims=" %%P in ('py -3.11 -c "import sys; print(sys.executable)" 2^>nul') do (
  set "CANDIDATE=%%P"
  call :TEST_BASE_PY
)
if defined BASE_PY exit /b 0

for /f "tokens=2,*" %%A in ('reg query "HKCU\Software\Python\PythonCore\3.11\InstallPath" /ve 2^>nul ^| findstr /I "REG_SZ"') do (
  set "CANDIDATE=%%B"
  if "!CANDIDATE:~-1!"=="\" (
    set "CANDIDATE=!CANDIDATE!python.exe"
  ) else (
    set "CANDIDATE=!CANDIDATE!\python.exe"
  )
  call :TEST_BASE_PY
)
if defined BASE_PY exit /b 0

for /f "tokens=2,*" %%A in ('reg query "HKLM\Software\Python\PythonCore\3.11\InstallPath" /ve 2^>nul ^| findstr /I "REG_SZ"') do (
  set "CANDIDATE=%%B"
  if "!CANDIDATE:~-1!"=="\" (
    set "CANDIDATE=!CANDIDATE!python.exe"
  ) else (
    set "CANDIDATE=!CANDIDATE!\python.exe"
  )
  call :TEST_BASE_PY
)
if defined BASE_PY exit /b 0

for /f "tokens=2,*" %%A in ('reg query "HKLM\Software\WOW6432Node\Python\PythonCore\3.11\InstallPath" /ve 2^>nul ^| findstr /I "REG_SZ"') do (
  set "CANDIDATE=%%B"
  if "!CANDIDATE:~-1!"=="\" (
    set "CANDIDATE=!CANDIDATE!python.exe"
  ) else (
    set "CANDIDATE=!CANDIDATE!\python.exe"
  )
  call :TEST_BASE_PY
)
if defined BASE_PY exit /b 0

set "CANDIDATE=%LOCALAPPDATA%\Programs\Python\Python311\python.exe"
call :TEST_BASE_PY
if defined BASE_PY exit /b 0

set "CANDIDATE=%ProgramFiles%\Python311\python.exe"
call :TEST_BASE_PY
if defined BASE_PY exit /b 0

set "CANDIDATE=%ProgramFiles(x86)%\Python311\python.exe"
call :TEST_BASE_PY
if defined BASE_PY exit /b 0

for /f "delims=" %%P in ('where python 2^>nul') do (
  set "CANDIDATE=%%P"
  call :TEST_BASE_PY
)
exit /b 0

:TEST_BASE_PY
if defined BASE_PY exit /b 0
if not defined CANDIDATE exit /b 1
if not exist "%CANDIDATE%" exit /b 1
"%CANDIDATE%" -B -c "import sys; raise SystemExit(0 if sys.version_info[:2] == (3, 11) else 1)" >nul 2>nul
if not errorlevel 1 set "BASE_PY=%CANDIDATE%"
exit /b 0

:CREATE_VENV
if not defined BASE_PY exit /b 1
if exist "%ENV_DIR%" rmdir /s /q "%ENV_DIR%"
if exist "%ENV_DIR%" exit /b 1
"%BASE_PY%" -m venv "%ENV_DIR%"
if errorlevel 1 exit /b 1
call :SELECT_ENV_PY
if not defined PY exit /b 1
"%PY%" -m ensurepip --upgrade >nul 2>nul
exit /b 0

:DETECT_STALE_PY311
set "STALE_PY311="
reg query "HKLM\Software\Microsoft\Windows\CurrentVersion\Uninstall\{B074012B-9B85-4049-BA01-A58A8C4C2236}" >nul 2>nul
if not errorlevel 1 set "STALE_PY311=1"
reg query "HKLM\Software\Microsoft\Windows\CurrentVersion\Uninstall\{C038789C-DCB5-42D3-8C51-3BC9DDB26B90}" >nul 2>nul
if not errorlevel 1 set "STALE_PY311=1"
reg query "HKLM\Software\Microsoft\Windows\CurrentVersion\Uninstall\{9AFDC691-40E5-4B15-835F-9A524AC4672C}" >nul 2>nul
if not errorlevel 1 set "STALE_PY311=1"
reg query "HKLM\Software\Microsoft\Windows\CurrentVersion\Uninstall\{1D653E80-09B2-40AB-9530-7DF12E632F8A}" >nul 2>nul
if not errorlevel 1 set "STALE_PY311=1"
reg query "HKLM\Software\Microsoft\Windows\CurrentVersion\Uninstall\{57AC2A86-EC99-4E9C-9FF9-15DAA88D1FAE}" >nul 2>nul
if not errorlevel 1 set "STALE_PY311=1"
exit /b 0
