#!/bin/bash
# VLAlert-X Phase 2.1 — SFT Qwen3-VL-4B on merged GPT-5 distilled CoT corpus.
#
# Reuses the existing per-frame CoT-belief trainer (training/VLA/train_cot_belief.py)
# but with:
#   - LoRA r=64 (was 32) for richer adaptation under the 1338-clip extension
#   - Lower LR (5e-5) since we're warm-starting from a checkpoint that has
#     already been SFT'd on the 2k-clip corpus
#   - Source manifest: data/cot_corpus_v2/vlalert_x_sft.jsonl (1338 records)
#   - Optional union with the existing perframe corpus
#
# Preconditions:
#   - data/cot_corpus_v2/vlalert_x_sft.jsonl   (run merge_gpt5_into_cot_manifest.py)
#   - checkpoints/VLA/qwen3vl4b_cot_belief_perframe/best/  (warm-start)
#   - models/Qwen3-VL-4B-Instruct/             (base VLM)
#
# Usage:
#   bash training/SFT/train_sft_x.sh              # full run (~3 GPU-hr)
#   bash training/SFT/train_sft_x.sh --debug      # 16-clip smoke
#   bash training/SFT/train_sft_x.sh --no_union   # GPT-5 corpus only

set -euo pipefail
cd "$(dirname "$0")/../.."
ROOT=$(pwd)

MODEL_PATH="${MODEL_PATH:-$ROOT/models/Qwen3-VL-4B-Instruct}"
GPT5_JSONL="${GPT5_JSONL:-$ROOT/data/cot_corpus_v2/vlalert_x_sft.jsonl}"
LEGACY_JSONL="${LEGACY_JSONL:-$ROOT/data/vla_cot_belief/train_perframe_union.jsonl}"
UNION_JSONL="${UNION_JSONL:-$ROOT/data/cot_corpus_v2/vlalert_x_train_union.jsonl}"

# Inference path for cot_belief_dataset; the dataset itself reads
# `video_path` from each record so VIDEO_DIR is a fallback only.
VIDEO_DIR="${VIDEO_DIR:-$ROOT/nexar-collision-prediction/train}"
RESUME="${RESUME:-$ROOT/checkpoints/VLA/qwen3vl4b_cot_belief_perframe/best}"
OUT_DIR="${OUT_DIR:-$ROOT/checkpoints/sft_x}"

EPOCHS="${EPOCHS:-2}"           # warm-start: 2 epoch enough on 1338 clips
BATCH_SIZE="${BATCH_SIZE:-1}"
GRAD_ACCUM="${GRAD_ACCUM:-8}"
LR="${LR:-5e-5}"                # lower than 1e-4 (warm-start, smaller corpus)
N_FRAMES="${N_FRAMES:-8}"
LORA_R="${LORA_R:-64}"          # ↑ from 32 to 64 (VLAlert-X plan)
MAX_LEN="${MAX_LEN:-4096}"
RESIZE_SHORT="${RESIZE_SHORT:-336}"

# ── flags ──
USE_UNION=1
DEBUG_ARGS=""
for arg in "$@"; do
    case "$arg" in
        --no_union) USE_UNION=0 ;;
        --debug)    DEBUG_ARGS="--max_samples 16 --epochs 1 --log_every 1" ;;
    esac
done

# ── build training manifest ──
mkdir -p "$(dirname "$UNION_JSONL")"
n_gpt5=$(wc -l < "$GPT5_JSONL")
if [[ $USE_UNION -eq 1 && -f "$LEGACY_JSONL" ]]; then
    cat "$GPT5_JSONL" "$LEGACY_JSONL" > "$UNION_JSONL"
    n_legacy=$(wc -l < "$LEGACY_JSONL")
    n_total=$(wc -l < "$UNION_JSONL")
    echo "[union] gpt5=$n_gpt5  legacy=$n_legacy  total=$n_total → $UNION_JSONL"
else
    cp "$GPT5_JSONL" "$UNION_JSONL"
    n_total=$(wc -l < "$UNION_JSONL")
    echo "[union] gpt5-only=$n_gpt5  → $UNION_JSONL"
fi

# ── pre-flight ──
for f in "$MODEL_PATH" "$UNION_JSONL"; do
    if [[ ! -e "$f" ]]; then
        echo "[FAIL] missing: $f" >&2
        exit 2
    fi
done

RESUME_ARG=""
if [[ -n "$RESUME" && -e "$RESUME/adapter_config.json" ]]; then
    RESUME_ARG="--resume $RESUME"
    echo "[resume] warm-start from $RESUME"
else
    echo "[resume] no warm-start — fresh LoRA init"
fi

mkdir -p "$OUT_DIR"
LOG_FILE="$ROOT/runs/vlalert_x/phase2_1_sft_$(date +%Y%m%d_%H%M%S).log"
mkdir -p "$(dirname "$LOG_FILE")"

echo "[config] EPOCHS=$EPOCHS  BATCH=$BATCH_SIZE  GRAD_ACCUM=$GRAD_ACCUM"
echo "         LR=$LR  LORA_R=$LORA_R  N_FRAMES=$N_FRAMES"
echo "         OUT_DIR=$OUT_DIR"
echo "         LOG_FILE=$LOG_FILE"

# IMPORTANT: route through the fast wrapper that applies the
# Conv3d→Linear patch (PR/qwen3vl_patch_embed_conv3d_slowdown.md).
# Skipping this gives ~17s/it (not the ~0.3s/it expected on RTX 5090).
python -m tools.run_train_cot_belief_fast \
    --model_name   "$MODEL_PATH" \
    --cot_jsonl    "$UNION_JSONL" \
    --video_dir    "$VIDEO_DIR" \
    --out_dir      "$OUT_DIR" \
    --epochs       "$EPOCHS" \
    --batch_size   "$BATCH_SIZE" \
    --grad_accum   "$GRAD_ACCUM" \
    --lr           "$LR" \
    --n_frames     "$N_FRAMES" \
    --lora_r       "$LORA_R" \
    --max_len      "$MAX_LEN" \
    --resize_short "$RESIZE_SHORT" \
    --per_frame \
    $RESUME_ARG \
    --save_every_epoch \
    $DEBUG_ARGS \
    2>&1 | tee "$LOG_FILE"

echo
echo "[done] checkpoint: $OUT_DIR/best/"
echo "[next] bash scripts/run_vlalert_x_pipeline.sh phase2_2_full"