#!/bin/bash
# VLAlert-X v2 Phase 1 — re-SFT Qwen3-VL-4B with BELIEF reasoning format.
#
# New prompt format (per frame):
#   <|BELIEF|> {per-frame reasoning text} </|BELIEF|> <|ACTION_i|>
#
# Two-stage LR schedule (sequential calls):
#   Stage 1A: lr=1e-4, 3 epochs, fresh LoRA r=128
#   Stage 1B: lr=2e-5, 2 epochs, warm-start from 1A best
# Total: 5 epochs, ~25 GPU-hr on RTX 5090.
#
# Usage:
#   bash training/SFT/train_sft_x_v2.sh             # full pipeline (1A + 1B)
#   bash training/SFT/train_sft_x_v2.sh smoke       # 50-sample smoke (5 min)
#   bash training/SFT/train_sft_x_v2.sh stage1a     # only 1A
#   bash training/SFT/train_sft_x_v2.sh stage1b     # only 1B (requires 1A done)
set -euo pipefail
cd "$(dirname "$0")/../.."

OUT_DIR="checkpoints/sft_x_v2"
TRAIN_JSONL="data/cot_corpus_v2/vlalert_x_perframe_v2_train.jsonl"
VAL_JSONL="data/cot_corpus_v2/vlalert_x_perframe_v2_val.jsonl"
mkdir -p logs "$OUT_DIR"

step="${1:-all}"

run_smoke() {
  echo "================================================================"
  echo "[smoke] 50 samples × 1 epoch (~5 min)"
  echo "================================================================"
  python -m training.VLA.train_cot_belief_v2 \
      --train_jsonl "$TRAIN_JSONL" \
      --val_jsonl   "$VAL_JSONL" \
      --out_dir     "${OUT_DIR}_smoke" \
      --epochs 1 \
      --batch_size 1 --grad_accum 2 \
      --lora_r 64 --lora_alpha 16 --lr 1e-4 \
      --max_samples 50 \
      --action_token_weight 2.0 \
      --log_every 5 2>&1 | tee logs/phase1_smoke.log
}

run_stage1a() {
  echo "================================================================"
  echo "[Stage 1A] full corpus × 3 epochs at lr=1e-4 (broad learning)"
  echo "  batch=2, grad_accum=2 (effective batch=4), LoRA r=128, action_w=2"
  echo "  Conv3d→Linear PR patch active (~17× per-step speedup)"
  echo "================================================================"
  python -m training.VLA.train_cot_belief_v2 \
      --train_jsonl "$TRAIN_JSONL" \
      --val_jsonl   "$VAL_JSONL" \
      --out_dir     "${OUT_DIR}/stage1a" \
      --epochs 3 \
      --batch_size 2 --grad_accum 2 \
      --lora_r 128 --lora_alpha 32 --lora_dropout 0.05 \
      --lr 1e-4 \
      --action_token_weight 2.0 \
      --save_every_epoch \
      --log_every 50 2>&1 | tee logs/phase1a_stage1a.log
}

run_stage1b() {
  echo "================================================================"
  echo "[Stage 1B] full corpus × 2 epochs at lr=2e-5 (fine-tune from 1A best)"
  echo "  batch=2, grad_accum=2, warm-start from Stage 1A"
  echo "================================================================"
  if [[ ! -d "${OUT_DIR}/stage1a/best" ]]; then
    echo "[FAIL] missing ${OUT_DIR}/stage1a/best — run stage1a first" >&2
    exit 1
  fi
  python -m training.VLA.train_cot_belief_v2 \
      --train_jsonl "$TRAIN_JSONL" \
      --val_jsonl   "$VAL_JSONL" \
      --out_dir     "${OUT_DIR}/stage1b" \
      --epochs 2 \
      --batch_size 2 --grad_accum 2 \
      --lora_r 128 --lora_alpha 32 --lora_dropout 0.05 \
      --lr 2e-5 \
      --action_token_weight 2.0 \
      --resume "${OUT_DIR}/stage1a/best" \
      --save_every_epoch \
      --log_every 50 2>&1 | tee logs/phase1b_stage1b.log
  # Promote stage1b/best as the final ckpt
  rm -rf "${OUT_DIR}/best"
  cp -r "${OUT_DIR}/stage1b/best" "${OUT_DIR}/best"
  echo "[done] final adapter -> ${OUT_DIR}/best"
}

case "$step" in
  smoke)    run_smoke ;;
  stage1a)  run_stage1a ;;
  stage1b)  run_stage1b ;;
  all)      run_stage1a && run_stage1b ;;
  *) echo "usage: $0 [smoke|stage1a|stage1b|all]" >&2; exit 2 ;;
esac