#!/bin/bash
# Anti-UAV Environment Rebuild Script
# 解决cuDNN兼容性问题，创建稳定GPU训练环境

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

echo "🔧 Anti-UAV 环境重建脚本"
echo "==============================="

# 1. 环境信息收集
echo "📊 当前环境分析..."
echo "Python: $(python --version)"
echo "Conda Environment: $CONDA_DEFAULT_ENV"
echo "GPU: $(nvidia-smi --query-gpu=name --format=csv,noheader)"

# 2. 创建兼容环境的requirements
echo "📝 生成兼容版本依赖列表..."
cat > /tmp/anti_uav_compatible_requirements.txt << 'EOF'
# Anti-UAV Compatible Dependencies
# PyTorch 2.3.0 + CUDA 12.1 (稳定组合)
torch==2.3.0+cu121
torchvision==0.18.0+cu121
torchaudio==2.3.0+cu121
--extra-index-url https://download.pytorch.org/whl/cu121

# Core ML Libraries (兼容版本)
ultralytics==8.2.0  # 稳定版本，避免最新版本问题
opencv-python==4.8.1.78
pillow==10.0.1
numpy==1.24.4
scipy==1.11.4

# Training Utilities
tqdm==4.66.1
matplotlib==3.7.5
seaborn==0.13.0
tensorboard==2.15.1

# Optional High-Performance Dependencies
optuna==3.4.0          # 超参数优化
wandb==0.16.0          # 实验追踪
albumentations==1.3.1  # 数据增强

# Development Tools
ipython==8.18.1
jupyter==1.0.0
pytest==7.4.3
EOF

echo "✅ 兼容版本依赖列表已生成"

# 3. 备份当前环境
echo "💾 备份当前环境..."
if command -v conda &> /dev/null; then
    conda env export > /tmp/backup_anti_uav_env_$(date +%Y%m%d_%H%M%S).yml
    echo "✅ Conda环境已备份到 /tmp/"
fi

# 4. 卸载冲突的PyTorch版本
echo "🗑️  卸载当前PyTorch版本..."
pip uninstall -y torch torchvision torchaudio ultralytics || true

# 5. 安装兼容版本
echo "📦 安装兼容版本依赖..."
pip install -r /tmp/anti_uav_compatible_requirements.txt

# 6. 验证安装
echo "🔍 验证新环境..."
python -c "
import torch
print(f'✅ PyTorch: {torch.__version__}')
print(f'✅ CUDA Available: {torch.cuda.is_available()}')
if torch.cuda.is_available():
    print(f'✅ CUDA Version: {torch.version.cuda}')
    print(f'✅ cuDNN Version: {torch.backends.cudnn.version()}')
    print(f'✅ GPU Count: {torch.cuda.device_count()}')
    print(f'✅ GPU Name: {torch.cuda.get_device_name(0)}')
else:
    print('❌ CUDA不可用')
"

# 7. 测试基本GPU功能
echo "🧪 GPU功能测试..."
python -c "
import torch
if torch.cuda.is_available():
    x = torch.randn(100, 100).cuda()
    y = torch.randn(100, 100).cuda()
    z = torch.matmul(x, y)
    print('✅ GPU矩阵运算测试通过')
    
    # 测试cuDNN
    from torch.backends import cudnn
    if cudnn.is_available():
        print('✅ cuDNN可用且功能正常')
    else:
        print('⚠️  cuDNN不可用')
else:
    print('❌ GPU功能测试失败')
"

echo "🎯 环境重建完成！"
echo "==============================="
echo "下一步: 运行GPU训练测试"