#!/bin/bash
# 训练前环境检查脚本

echo "======================================"
echo "BEVFusion 训练环境检查"
echo "======================================"
echo ""

# 1. 检查Python和PyTorch
echo "[1/8] 检查Python环境..."
if command -v python3 &> /dev/null; then
    PYTHON_VERSION=$(python3 --version 2>&1)
    echo "  ✓ Python: $PYTHON_VERSION"
else
    echo "  ✗ Python3未找到"
    exit 1
fi

# 2. 检查GPU
echo ""
echo "[2/8] 检查GPU..."
if command -v nvidia-smi &> /dev/null; then
    GPU_COUNT=$(nvidia-smi --query-gpu=name --format=csv,noheader | wc -l)
    echo "  ✓ 可用GPU数量: $GPU_COUNT"
    nvidia-smi --query-gpu=name,memory.total --format=csv,noheader | nl -v 0 | sed 's/^/    GPU /'
else
    echo "  ✗ nvidia-smi未找到"
    exit 1
fi

# 3. 检查torchpack
echo ""
echo "[3/8] 检查torchpack..."
if command -v torchpack &> /dev/null; then
    echo "  ✓ torchpack已安装"
else
    echo "  ✗ torchpack未安装"
    echo "    安装命令: pip install torchpack"
    exit 1
fi

# 4. 检查预训练模型
echo ""
echo "[4/8] 检查预训练模型..."
PRETRAIN_DIR="pretrained"
REQUIRED_MODELS=(
    "swint-nuimages-pretrained.pth"
    "lidar-only-det.pth"
)

for model in "${REQUIRED_MODELS[@]}"; do
    if [ -f "$PRETRAIN_DIR/$model" ]; then
        SIZE=$(du -h "$PRETRAIN_DIR/$model" | cut -f1)
        echo "  ✓ $model ($SIZE)"
    else
        echo "  ✗ $model 未找到"
    fi
done

# 5. 检查数据集
echo ""
echo "[5/8] 检查nuScenes数据集..."
DATA_ROOT="data/nuscenes"
REQUIRED_DATA=(
    "nuscenes_infos_train.pkl"
    "nuscenes_infos_val.pkl"
    "samples"
    "sweeps"
)

if [ -d "$DATA_ROOT" ]; then
    echo "  ✓ 数据集目录存在: $DATA_ROOT"
    for data in "${REQUIRED_DATA[@]}"; do
        if [ -e "$DATA_ROOT/$data" ]; then
            echo "    ✓ $data"
        else
            echo "    ✗ $data 未找到"
        fi
    done
else
    echo "  ✗ 数据集目录不存在: $DATA_ROOT"
fi

# 6. 检查配置文件
echo ""
echo "[6/8] 检查多任务配置文件..."
CONFIG_FILE="configs/nuscenes/multitask/fusion-det-seg-swint.yaml"
if [ -f "$CONFIG_FILE" ]; then
    echo "  ✓ 配置文件存在: $CONFIG_FILE"
else
    echo "  ✗ 配置文件未找到: $CONFIG_FILE"
    exit 1
fi

# 7. 检查自定义算子
echo ""
echo "[7/8] 检查自定义CUDA算子..."
python3 -c "
try:
    from mmdet3d.ops import bev_pool_v2
    print('  ✓ BEV Pool算子')
except Exception as e:
    print('  ✗ BEV Pool算子:', str(e))
    
try:
    from mmdet3d.ops import Voxelization
    print('  ✓ Voxelization算子')
except Exception as e:
    print('  ✗ Voxelization算子:', str(e))

try:
    from mmdet3d.ops.spconv import SparseConv3d
    print('  ✓ Sparse Conv算子')
except Exception as e:
    print('  ✗ Sparse Conv算子:', str(e))
" 2>&1

# 8. 检查磁盘空间
echo ""
echo "[8/8] 检查磁盘空间..."
AVAILABLE_SPACE=$(df -h . | awk 'NR==2 {print $4}')
echo "  可用空间: $AVAILABLE_SPACE"

echo ""
echo "======================================"
echo "环境检查完成"
echo "======================================"