#!/usr/bin/env bash
# SFT v3: 提升 Belief Vector 质量
#
# 相比 v2 的改动（每项都有明确理由）：
#
#   1. belief_strategy: mean_pool → attention_pool
#      理由: mean_pool 对所有 token 平均权重相同；attention_pool 让模型自动学习
#            "哪些视觉 token 对危险预测最重要"，得到更具区分度的 belief vector。
#
#   2. curriculum: 去掉 --no_curriculum（重新启用）
#      理由: v2 关闭了课程学习。启用后训练从简单样本（明确碰撞/明确安全）开始，
#            逐步引入难例，防止模型早期过拟合噪声标签。
#
#   3. nll_weight: 0.5 → 0.3
#      理由: NLL 是 VLM 语言生成 loss（让模型输出正确文字）。
#            对 policy 训练，我们需要的是 belief vector 的质量，
#            而不是文字生成质量。降低 nll_weight 让 hazard + TTA heads
#            对 LoRA 更新有更大的梯度影响力。
#
#   4. vlm_lr_multiplier: 0.1 → 0.05
#      理由: VLM（LoRA 部分）更新更保守，防止破坏 Qwen2.5 的预训练知识。
#            头部 (belief_aggregator, hazard_head, tta_head) 用正常 lr 学习。
#
#   5. num_epochs: 10 → 12，tta_head_lr: 1e-3 → 2e-3
#      理由: attention_pool 引入了额外参数，需要更多训练步数收敛。
#            tta_head 负责时间预测，给更高 lr 让它更快适应任务。
#
# 期望结果（相比 v2 的 belief vector）：
#   - 更好的 hazard 判断 (AUC ↑ ~0.02)
#   - 更准确的 TTA 估计 (RMSE ↓ ~0.3s)
#   - 下游 policy 分 +0.02~0.04 (policy_score 从 ~0.77 → ~0.80+)
#
# 前提：v2 checkpoint 用于初始化（热启动，不是从头训练）
#
# 运行时间估计：~8-10 小时（单 GPU）
#
# 完整训练链（SFT v3 之后需要）：
#   Step A: bash training/SFT/train_sft_v3.sh        # 本脚本 ~8-10h
#   Step B: python -m training.Policy.cache_beliefs   # 重新生成 belief 缓存 ~1h
#   Step C: bash training/Policy/train_policy_v3.sh   # 重新训练 policy ~15min
#           (改 SFT_CHECKPOINT → checkpoints/SFT/sft_v3/best)
#
# 用法：
#   bash training/SFT/train_sft_v3.sh             # 完整训练
#   bash training/SFT/train_sft_v3.sh --debug     # 调试 (~5 min)
set -euo pipefail

ROOT=PROJECT_ROOT
MANIFEST_DIR="$ROOT/data/sft_manifests"
PRETRAINED_LORA="$ROOT/checkpoints/SFT/sft_v2/best"   # 从 v2 热启动
MODEL_PATH="$ROOT/models/Qwen2.5-VL-3B-Instruct"
OUTPUT_DIR="$ROOT/checkpoints/SFT"
EXPERIMENT="sft_v3"

MAX_PIXELS=602112   # 768*28*28 — 与 v2 相同

BATCH_SIZE=2
GRAD_ACCUM=4        # effective batch = 8，与 v2 相同

DEBUG_FLAGS=""
if [[ "${1:-}" == "--debug" ]]; then
    DEBUG_FLAGS="--debug --debug_samples 64"
    EXPERIMENT="sft_v3_debug"
    BATCH_SIZE=2
    GRAD_ACCUM=4
    echo "=== DEBUG / SMOKE TEST MODE ==="
fi

# ── 确认 v2 checkpoint 存在 ──────────────────────────────────────────────────
if [[ ! -d "$PRETRAINED_LORA" ]]; then
    echo "⚠ SFT v2 checkpoint 不存在: $PRETRAINED_LORA"
    echo "  → 将从 pretrain_v2/stage_b/best_model 冷启动"
    PRETRAINED_LORA="$ROOT/checkpoints/pretrain_v2/stage_b/best_model"
fi

# ── 确认 manifests 存在 ──────────────────────────────────────────────────────
if [[ ! -f "$MANIFEST_DIR/nexar_train.json" ]]; then
    echo "Manifests not found — generating..."
    python -m training.SFT.make_split_manifest \
        --nexar_root "$ROOT/NEXAR_COLLISION/dataset" \
        --dada_root  "$ROOT/DADA-2000" \
        --out_dir    "$MANIFEST_DIR"
fi

echo ""
echo "=== SFT v3 训练 ==="
echo "  改动汇总 vs v2:"
echo "    belief_strategy  : mean_pool → attention_pool"
echo "    curriculum       : False → True"
echo "    nll_weight       : 0.5 → 0.3"
echo "    vlm_lr_multiplier: 0.1 → 0.05"
echo "    num_epochs       : 10  → 12"
echo "    tta_head_lr      : 1e-3 → 2e-3"
echo ""
echo "  热启动自: $PRETRAINED_LORA"
echo "  输出至  : $OUTPUT_DIR/$EXPERIMENT"
echo ""

python -m training.SFT.trainer \
    --manifest_dir      "$MANIFEST_DIR"     \
    --model_name        "$MODEL_PATH"       \
    --pretrained_lora   "$PRETRAINED_LORA"  \
    --output_dir        "$OUTPUT_DIR"       \
    --experiment_name   "$EXPERIMENT"       \
    --belief_strategy   attention_pool      \
    --num_epochs        12                  \
    --batch_size        $BATCH_SIZE         \
    --gradient_accumulation_steps $GRAD_ACCUM \
    --learning_rate     1e-4                \
    --tta_head_lr       2e-3                \
    --vlm_lr_multiplier 0.05                \
    --weight_decay      0.01                \
    --max_grad_norm     1.0                 \
    --nll_weight        0.3                 \
    --max_pixels        $MAX_PIXELS         \
    --no_auto_resume                        \
    --use_wandb                             \
    $DEBUG_FLAGS

echo ""
echo "✅ SFT v3 训练完成"
echo "   Checkpoint: $OUTPUT_DIR/$EXPERIMENT/best"
echo ""
echo "下一步 — 重新生成 belief 缓存:"
echo "  python -m training.Policy.cache_beliefs \\"
echo "      --sft_checkpoint $OUTPUT_DIR/$EXPERIMENT/best \\"
echo "      --label_dir      data/policy_labels \\"
echo "      --output_dir     data/belief_cache_v3"
echo ""
echo "然后用新缓存重训 policy:"
echo "  修改 train_policy_v3.sh 中 CACHE_DIR=data/belief_cache_v3"
echo "  并将 SFT_CHECKPOINT 改为 $OUTPUT_DIR/$EXPERIMENT/best"