#!/bin/bash
#
# 增强版多任务训练脚本
# 训练BEVFusion with EnhancedBEVSegmentationHead
#
# 预期性能：
# - 检测mAP: 65-68%
# - 分割mIoU: 60-65% (vs 当前36%)
#
# 改进点：
# 1. ASPP多尺度特征
# 2. 通道+空间注意力
# 3. 深层解码器(4层)
# 4. Focal Loss修复(alpha=0.25)
# 5. Dice Loss混合
# 6. 类别权重平衡
# 7. Deep supervision
# 8. 分割损失权重3.0x

set -e

echo "=========================================="
echo "BEVFusion 增强版多任务训练"
echo "任务: 3D检测 + BEV分割（增强版）"
echo "=========================================="
echo ""

# 配置
CONFIG="configs/nuscenes/multitask/fusion-det-seg-swint-enhanced.yaml"
NUM_GPUS=8
WORK_DIR="runs/enhanced-multitask-$(date +%Y%m%d-%H%M%S)"
LOG_FILE="training_enhanced_multitask_$(date +%Y%m%d_%H%M%S).log"

echo "配置信息："
echo "  配置文件: ${CONFIG}"
echo "  GPU数量: ${NUM_GPUS}"
echo "  工作目录: ${WORK_DIR}"
echo "  日志文件: ${LOG_FILE}"
echo ""

# 检查配置文件
if [ ! -f "${CONFIG}" ]; then
    echo "错误: 配置文件不存在: ${CONFIG}"
    exit 1
fi

# 检查预训练权重
PRETRAINED_SWINT="pretrained/swint-nuimages-pretrained.pth"
PRETRAINED_LIDAR="pretrained/lidar-only-det.pth"

if [ ! -f "${PRETRAINED_SWINT}" ]; then
    echo "警告: SwinT预训练权重不存在: ${PRETRAINED_SWINT}"
    echo "将从头开始训练camera encoder（不推荐）"
fi

if [ ! -f "${PRETRAINED_LIDAR}" ]; then
    echo "警告: LiDAR预训练权重不存在: ${PRETRAINED_LIDAR}"
    echo "将从头开始训练LiDAR encoder（不推荐）"
fi

echo ""
echo "改进点："
echo "  ✅ ASPP多尺度特征提取"
echo "  ✅ 通道+空间注意力机制"
echo "  ✅ 深层解码器(4层)"
echo "  ✅ Focal Loss修复(alpha=0.25)"
echo "  ✅ Dice Loss混合(weight=0.5)"
echo "  ✅ 类别权重平衡(stop_line×4)"
echo "  ✅ Deep supervision"
echo "  ✅ 分割损失权重×3"
echo ""

echo "预期性能:"
echo "  - 检测mAP: 65-68%"
echo "  - 分割mIoU: 60-65%"
echo "  - 训练时间: ~4-5天"
echo ""

echo "=========================================="
echo "开始训练..."
echo "=========================================="
echo ""

# 构建训练命令
TRAIN_CMD="torchpack dist-run -np ${NUM_GPUS} python tools/train.py ${CONFIG}"

# 添加预训练权重
if [ -f "${PRETRAINED_SWINT}" ]; then
    TRAIN_CMD="${TRAIN_CMD} --model.encoders.camera.backbone.init_cfg.checkpoint ${PRETRAINED_SWINT}"
fi

if [ -f "${PRETRAINED_LIDAR}" ]; then
    TRAIN_CMD="${TRAIN_CMD} --load_from ${PRETRAINED_LIDAR}"
fi

# 添加工作目录
TRAIN_CMD="${TRAIN_CMD} --run-dir ${WORK_DIR}"

# 输出训练命令
echo "训练命令:"
echo "${TRAIN_CMD}"
echo ""

# 执行训练（输出到文件和终端）
eval "${TRAIN_CMD}" 2>&1 | tee "${LOG_FILE}"

echo ""
echo "=========================================="
echo "训练完成！"
echo "=========================================="
echo ""
echo "结果保存在:"
echo "  工作目录: ${WORK_DIR}"
echo "  日志文件: ${LOG_FILE}"
echo ""
echo "评估命令:"
echo "torchpack dist-run -np ${NUM_GPUS} python tools/test.py \\"
echo "  ${CONFIG} \\"
echo "  ${WORK_DIR}/latest.pth \\"
echo "  --eval bbox map"
echo ""