@echo off
chcp 65001 >nul

REM 检查.venv目录是否存在
if exist .venv (
    echo 虚拟环境已存在。
) else (
    echo 正在创建虚拟环境...
    uv venv
    if errorlevel 1 (
        echo 创建虚拟环境失败。
        pause
        exit /b 1
    )
    echo.
    echo 虚拟环境已创建。请激活它后重新运行此脚本。
    echo   - Windows (CMD): .\.venv\Scripts\activate
    echo   - Windows (PowerShell): .\.venv\Scripts\Activate.ps1
    echo   - Linux/macOS: source .venv/bin/activate
    echo.
    pause
    exit /b 0
)

echo 安装PyTorch GPU版本（通过官方源）
echo ===================================
echo.

echo 请选择CUDA版本:
echo 1. CUDA 12.8 (适用于较新的NVIDIA GPU)
echo 2. CUDA 12.6 (适用于较旧的NVIDIA GPU)
echo 3. CUDA 11.8 (适用于较旧的NVIDIA GPU)
echo.

set /p choice=请输入选项 (1/2): 

if "%choice%"=="1" (
    echo 正在安装PyTorch CUDA 12.8版本...
    uv pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu128 -i https://pypi.tuna.tsinghua.edu.cn/simple
) 
else if "%choice%"=="2" (
    echo 正在安装PyTorch CUDA 12.6版本...
    uv pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu126 -i https://pypi.tuna.tsinghua.edu.cn/simple
)
else if "%choice%"=="3" (
    echo 正在安装PyTorch CUDA 11.8版本...
    uv pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118 -i https://pypi.tuna.tsinghua.edu.cn/simple
) else (
    echo 无效的选项！
    goto end
)

echo.
echo 验证PyTorch GPU支持状态...
python -c "import torch; print('CUDA是否可用:', torch.cuda.is_available()); print('PyTorch版本:', torch.__version__); print('CUDA版本:', torch.version.cuda if torch.cuda.is_available() else '无')"

:end
pause