#!/bin/bash

# =====================================================================
# Anti-UAV PyTorch 环境完整修复脚本
# 解决CUDA版本冲突和配置不一致问题
# =====================================================================

set -e  # 遇到错误立即退出

ENV_NAME="anti_uav_complete"
echo "🔧 开始修复 $ENV_NAME 环境..."

# 检查当前环境
if [[ "$CONDA_DEFAULT_ENV" != "$ENV_NAME" ]]; then
    echo "❌ 请先激活环境: conda activate $ENV_NAME"
    exit 1
fi

echo "📊 当前版本检查..."
echo "=== CUDA驱动版本 ==="
nvidia-smi | grep "CUDA Version"
echo "=== CUDA工具链版本 ==="
nvcc --version | grep "release"
echo "=== PyTorch版本 ==="
python -c "import torch; print(f'PyTorch: {torch.__version__}'); print(f'CUDA: {torch.version.cuda}')"

echo ""
echo "🎯 问题诊断："
echo "  - 驱动支持: CUDA 12.4"
echo "  - 工具链版本: CUDA 11.5" 
echo "  - PyTorch版本: CUDA 12.1"
echo "  - cuDNN冲突: 版本不匹配"
echo ""

read -p "是否继续修复？这将重新安装PyTorch (y/N): " confirm
if [[ $confirm != [yY] ]]; then
    echo "修复已取消"
    exit 0
fi

echo "🔄 步骤1: 卸载冲突的PyTorch组件..."
pip uninstall -y torch torchvision torchaudio pytorch-cuda || true

echo "🔄 步骤2: 清理CUDA缓存..."
python -c "import torch; torch.cuda.empty_cache()" 2>/dev/null || true

echo "🔄 步骤3: 安装兼容的PyTorch版本 (CUDA 11.8)..."
# 使用CUDA 11.8版本，这个版本与CUDA 11.5更兼容
pip install torch==2.3.0+cu118 torchvision==0.18.0+cu118 torchaudio==2.3.0+cu118 \
    --index-url https://download.pytorch.org/whl/cu118

echo "🔄 步骤4: 验证PyTorch安装..."
python -c "
import torch
print(f'PyTorch: {torch.__version__}')
print(f'CUDA available: {torch.cuda.is_available()}')
print(f'CUDA version: {torch.version.cuda}')
if torch.cuda.is_available():
    print(f'GPU: {torch.cuda.get_device_name()}')
    print('GPU测试...', end=' ')
    x = torch.randn(100, 100).cuda()
    y = torch.matmul(x, x)
    print('✅ 成功')
else:
    print('❌ CUDA不可用')
"

echo "🔄 步骤5: 重新安装Ultralytics..."
pip install ultralytics==8.2.0

echo "🔄 步骤6: 验证训练环境..."
python -c "
from ultralytics import YOLO
import torch
print('正在加载YOLOv8模型...', end=' ')
model = YOLO('yolov8n.pt')
print('✅ 成功')
print(f'设备: {\"cuda\" if torch.cuda.is_available() else \"cpu\"}')
"

echo ""
echo "✅ 环境修复完成！"
echo ""
echo "📋 修复后版本信息："
python -c "
import torch
from ultralytics import YOLO
print(f'PyTorch: {torch.__version__}')
print(f'CUDA: {torch.version.cuda}')
print(f'Ultralytics: 8.2.0')
print(f'GPU可用: {torch.cuda.is_available()}')
"

echo ""
echo "🎯 现在可以开始训练："
echo "  python tools/train_det.py --data configs/det/antiuav_test_small.yaml --epochs 10 --batch 8"