#!/bin/bash
# BEVFusion 多任务训练脚本（检测 + 分割）

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

echo "======================================"
echo "BEVFusion 多任务训练脚本"
echo "任务：3D目标检测 + BEV地图分割"
echo "======================================"
echo ""

# 配置参数
NUM_GPUS=8
CONFIG="configs/nuscenes/multitask/fusion-det-seg-swint.yaml"
BACKBONE_PRETRAIN="pretrained/swint-nuimages-pretrained.pth"
LIDAR_PRETRAIN="pretrained/lidar-only-det.pth"

# 检查预训练模型
echo "检查预训练模型..."
if [ ! -f "$BACKBONE_PRETRAIN" ]; then
    echo "错误：找不到 $BACKBONE_PRETRAIN"
    echo "请先下载预训练模型"
    exit 1
fi

if [ ! -f "$LIDAR_PRETRAIN" ]; then
    echo "错误：找不到 $LIDAR_PRETRAIN"
    echo "请先下载预训练模型"
    exit 1
fi
echo "✓ 预训练模型检查通过"
echo ""

# 检查配置文件
echo "检查配置文件..."
if [ ! -f "$CONFIG" ]; then
    echo "错误：找不到配置文件 $CONFIG"
    exit 1
fi
echo "✓ 配置文件检查通过"
echo ""

# 检查GPU
echo "检查GPU..."
GPU_COUNT=$(nvidia-smi --query-gpu=name --format=csv,noheader | wc -l)
echo "可用GPU数量: $GPU_COUNT"
if [ $GPU_COUNT -lt $NUM_GPUS ]; then
    echo "警告：可用GPU数量($GPU_COUNT)小于配置数量($NUM_GPUS)"
    echo "将使用 $GPU_COUNT 个GPU"
    NUM_GPUS=$GPU_COUNT
fi
echo ""

# 显示训练配置
echo "======================================"
echo "训练配置："
echo "  配置文件: $CONFIG"
echo "  GPU数量: $NUM_GPUS"
echo "  Camera Backbone预训练: $BACKBONE_PRETRAIN"
echo "  LiDAR预训练: $LIDAR_PRETRAIN"
echo "  任务: 3D检测 + BEV分割"
echo "======================================"
echo ""

# 询问确认
read -p "是否开始训练？(y/n) " -n 1 -r
echo
if [[ ! $REPLY =~ ^[Yy]$ ]]; then
    echo "训练已取消"
    exit 0
fi

# 开始训练
echo ""
echo "开始训练..."
echo ""

torchpack dist-run -np $NUM_GPUS python tools/train.py \
  $CONFIG \
  --model.encoders.camera.backbone.init_cfg.checkpoint $BACKBONE_PRETRAIN \
  --load_from $LIDAR_PRETRAIN

echo ""
echo "======================================"
echo "训练完成！"
echo "======================================"