#!/usr/bin/env bash
# SFT v2: dual-head (hazard + TTA), manifest-based, initialized from pretrain_v2 stage_b.
#
# GPU tuning: batch_size=4, grad_accum=2 → effective batch=8 (same as before).
# max_pixels=401408 (512*28*28) reduces per-sample VRAM vs default 768*28*28,
# allowing larger batch without OOM. If OOM, try batch_size=2 grad_accum=4.
#
# Usage:
#   bash training/SFT/train_sft_v2.sh            # full training
#   bash training/SFT/train_sft_v2.sh --debug    # smoke test (~300 steps)
set -euo pipefail

ROOT=PROJECT_ROOT
MANIFEST_DIR="$ROOT/data/sft_manifests"
PRETRAINED_LORA="$ROOT/checkpoints/pretrain_v2/stage_b/best_model"
MODEL_PATH="$ROOT/models/Qwen2.5-VL-3B-Instruct"
OUTPUT_DIR="$ROOT/checkpoints/SFT"
EXPERIMENT="sft_v2"

MAX_PIXELS=602112  # 768*28*28, Qwen2.5-VL default

BATCH_SIZE=2
GRAD_ACCUM=4   # effective batch = BATCH_SIZE * GRAD_ACCUM = 8

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

# ── Step 0: ensure manifests exist ──────────────────────────────────────────
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

# ── Step 1: dataset sanity check ────────────────────────────────────────────
echo "Running dataset sanity check..."
python -m training.SFT.sanity_check \
    --manifest_dir    "$MANIFEST_DIR" \
    --model_name      "$MODEL_PATH"   \
    --pretrained_lora "$PRETRAINED_LORA" \
    --skip_model

echo "Starting SFT v2 training..."
echo "  Manifests         : $MANIFEST_DIR"
echo "  Pretrained LoRA   : $PRETRAINED_LORA"
echo "  Output            : $OUTPUT_DIR/$EXPERIMENT"
echo "  batch_size        : $BATCH_SIZE  (grad_accum=$GRAD_ACCUM, eff_batch=$((BATCH_SIZE*GRAD_ACCUM)))"
echo "  max_pixels        : $MAX_PIXELS"

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