@echo off
chcp 65001 >nul
echo ============================================
echo   GTX 1660 TensorRT 优化版本安装脚本
echo ============================================
echo.

echo [1/5] 检查 Python 环境...
python --version
if errorlevel 1 (
    echo.
    echo ❌ 错误：未找到 Python！
    echo 请先安装 Python 3.8-3.11，并确保添加到 PATH
    pause
    exit /b 1
)
echo ✓ Python 已安装
echo.

echo [2/5] 升级 pip...
python -m pip install --upgrade pip
echo.

echo [3/5] 安装基础依赖...
pip install numpy opencv-python matplotlib kornia>=0.6.11
if errorlevel 1 (
    echo.
    echo ❌ 基础依赖安装失败！
    pause
    exit /b 1
)
echo ✓ 基础依赖安装完成
echo.

echo [4/5] 安装 PyTorch (CUDA 12.1)...
echo 正在卸载旧版本（如果有）...
pip uninstall -y torch torchvision torchaudio 2>nul
echo 正在安装 PyTorch 2.5.1 (CUDA 12.1)...
echo 这可能需要几分钟，请耐心等待...
pip install torch==2.5.1 torchvision==0.20.1 torchaudio==2.5.1 --index-url https://download.pytorch.org/whl/cu121
if errorlevel 1 (
    echo.
    echo ❌ PyTorch 安装失败！
    echo 请检查网络连接或手动安装
    pause
    exit /b 1
)
echo ✓ PyTorch 安装完成
echo.

echo [5/5] 安装 TensorRT (torch-tensorrt)...
pip install torch-tensorrt==2.5.0
if errorlevel 1 (
    echo.
    echo ⚠️  TensorRT 安装失败（可选，不影响基础使用）
    echo.
    echo 如果安装失败，你可以：
    echo 1. 稍后手动安装：pip install torch-tensorrt==2.5.0
    echo 2. 或使用 PyTorch FP16（不使用 TensorRT）
    echo 3. 或使用 ONNX → TensorRT 方案（见文档）
) else (
    echo ✓ TensorRT 安装完成
)
echo.

echo ============================================
echo   验证安装...
echo ============================================
echo.

echo 验证 PyTorch 和 CUDA...
python -c "import torch; print(f'✓ PyTorch: {torch.__version__}'); print(f'✓ CUDA available: {torch.cuda.is_available()}'); print(f'✓ CUDA version: {torch.version.cuda}'); gpu_name = torch.cuda.get_device_name(0) if torch.cuda.is_available() else 'N/A'; print(f'✓ GPU: {gpu_name}')"
echo.

echo 验证 TensorRT...
python -c "import torch_tensorrt; print('✓ TensorRT: 已安装')" 2>nul || echo "⚠️  TensorRT: 未安装（可选，不影响基础使用）"
echo.

echo ============================================
echo   安装完成！
echo ============================================
echo.
echo 现在你可以运行项目了：
echo.
echo 基础运行（不使用 TensorRT）：
echo   python demo_lightglue_camera_position_async.py --input "udp://0.0.0.0:12346" --max_keypoints 128 --use_fp16 --show_fps
echo.
echo 使用 TensorRT 加速：
echo   python demo_lightglue_camera_position_async.py --input "udp://0.0.0.0:12346" --max_keypoints 128 --use_fp16 --use_tensorrt --tensorrt_precision fp16 --show_fps
echo.
pause