#!/bin/bash
# Per-frame POMDP SFT on Qwen3-VL-4B (warm-start from existing clip-level ckpt).
#
# Preconditions:
#   - data/vla_cot_belief/train500_perframe.jsonl        (Nexar per-frame targets)
#   - data/vla_cot/dota_val_perframe.jsonl               (optional DoTA per-frame)
#   - checkpoints/VLA/qwen3vl4b_cot_belief/best          (clip-level warm-start)
#
# Usage:
#   bash training/VLA/train_cot_belief_perframe.sh              # full run
#   bash training/VLA/train_cot_belief_perframe.sh --debug      # smoke 16 clips
#
# GPU budget (5090 32GB): batch_size=1, grad_accum=8, max_len=4096 (per-frame
# target adds ~80 tokens, safely under 4k).

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

MODEL_PATH="${MODEL_PATH:-$(pwd)/models/Qwen3-VL-4B-Instruct}"
NEXAR_JSONL="${NEXAR_JSONL:-$(pwd)/data/vla_cot_belief/train500_perframe.jsonl}"
DOTA_JSONL="${DOTA_JSONL:-$(pwd)/data/vla_cot/dota_train_perframe.jsonl}"
UNION_JSONL="${UNION_JSONL:-$(pwd)/data/vla_cot_belief/train_perframe_union.jsonl}"
VIDEO_DIR="${VIDEO_DIR:-$(pwd)/nexar-collision-prediction/train}"
RESUME="${RESUME:-$(pwd)/checkpoints/VLA/qwen3vl4b_cot_belief/best}"
OUT_DIR="${OUT_DIR:-$(pwd)/checkpoints/VLA/qwen3vl4b_cot_belief_perframe}"

EPOCHS="${EPOCHS:-3}"
BATCH_SIZE="${BATCH_SIZE:-1}"
GRAD_ACCUM="${GRAD_ACCUM:-8}"
LR="${LR:-1e-4}"             # lower lr when warm-starting
N_FRAMES="${N_FRAMES:-8}"
LORA_R="${LORA_R:-32}"
MAX_LEN="${MAX_LEN:-4096}"
RESIZE_SHORT="${RESIZE_SHORT:-336}"

# Build union JSONL on the fly (Nexar + DoTA if DoTA file exists).
mkdir -p "$(dirname "$UNION_JSONL")"
if [[ -f "$DOTA_JSONL" ]]; then
    cat "$NEXAR_JSONL" "$DOTA_JSONL" > "$UNION_JSONL"
    n_nexar=$(wc -l < "$NEXAR_JSONL")
    n_dota=$(wc -l < "$DOTA_JSONL")
    n_total=$(wc -l < "$UNION_JSONL")
    echo "[union] nexar=$n_nexar  dota=$n_dota  total=$n_total -> $UNION_JSONL"
else
    cp "$NEXAR_JSONL" "$UNION_JSONL"
    n_total=$(wc -l < "$UNION_JSONL")
    echo "[union] DoTA JSONL not found; using Nexar-only ($n_total clips) -> $UNION_JSONL"
fi

for f in "$MODEL_PATH" "$UNION_JSONL" "$VIDEO_DIR"; 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

DEBUG_ARGS=""
if [[ "${1:-}" == "--debug" ]]; then
    DEBUG_ARGS="--max_samples 16 --epochs 1 --log_every 1"
    echo "[smoke] debug: 16 samples × 1 epoch"
fi

python -m training.VLA.train_cot_belief \
    --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