#!/bin/bash
# AR (Autoregressive) version for VILA WebDataset training
# Usage: bash train_vila_webds_ar.sh
# Or: DATASET_WDSMETA="sharegpt4v_sft" bash train_vila_webds_ar.sh

set -euxo pipefail

# Resolve project root and change to it immediately
PROJECT_ROOT="/lustre/fs11/portfolios/nvr/projects/nvr_elm_llm/users/chengyuew/workspace/shiyi/LMFlow"
cd "${PROJECT_ROOT}"

echo "=========================================="
echo "AR Qwen2.5-VL VLM Training (VILA WebDataset)"
echo "=========================================="
echo "Node: $(hostname)"
echo "Date: $(date)"
echo "Working Directory: $(pwd)"
echo ""

echo "Python: $(python3 --version)"
echo "PyTorch: $(python3 -c 'import torch; print(torch.__version__)')"
echo "CUDA available: $(python3 -c 'import torch; print(torch.cuda.is_available())')"
echo "CUDA device count: $(python3 -c 'import torch; print(torch.cuda.device_count())')"
echo "nvidia-smi:"
nvidia-smi || true
echo ""

# Set environment variables
export OMP_NUM_THREADS="${OMP_NUM_THREADS:-16}"
export PYTORCH_CUDA_ALLOC_CONF="expandable_segments:True,max_split_size_mb:64,garbage_collection_threshold:0.9"
export TORCHDYNAMO_DISABLE=1
export TORCH_COMPILE=0
export TOKENIZERS_PARALLELISM=false

# Force use of local src/ code instead of installed package
export PYTHONPATH="${PROJECT_ROOT}/src:${PYTHONPATH:-}"
echo "PYTHONPATH=${PYTHONPATH}"

# Set VILA dataset registry (adjust based on available YAML files in VILA/llava/data/registry/datasets/)
echo "VILA_DATASETS=${VILA_DATASETS}"

# Path mapping is handled by vila_path_patcher.py
echo "ℹ️  VILA path remapping will be handled by Python monkey-patching (vila_path_patcher.py)"
echo ""

# Set dataset backend to webdataset_multi_modal
# DATASET_BACKEND — loaded from cluster_env.sh
# Qwen2.5-VL processor dynamic resolution controls
export MIN_PIXELS="${MIN_PIXELS:-100352}"
export MAX_PIXELS="${MAX_PIXELS:-100352}"
export RESIZED_HEIGHT="${RESIZED_HEIGHT:-}"
export RESIZED_WIDTH="${RESIZED_WIDTH:-}"

# Configurable parameters
# AR model path (non-MDM)
MODEL_PATH="${MODEL_PATH:-/lustre/fs11/portfolios/nvr/projects/nvr_elm_llm/users/chengyuew/workspace/shiyi/LMFlow/src/models/Qwen2.5-VL-3B-Instruct}"
TOKENIZER_NAME="${TOKENIZER_NAME:-Qwen/Qwen2.5-VL-3B-Instruct}"

# VILA mixture name or .wdsmeta.json file
DATASET_WDSMETA="${DATASET_WDSMETA:-sharegpt4v_sft}"

EPOCHS="${EPOCHS:-1}"
BLOCK_SIZE="${BLOCK_SIZE:-2048}"
BATCH="${BATCH:-1}"
GRAD_ACC="${GRAD_ACC:-8}"
LR="${LR:-2e-5}"
WARMUP_RATIO="${WARMUP_RATIO:-0.03}"
BF16="${BF16:-1}"
GRADIENT_CHECKPOINTING="${GRADIENT_CHECKPOINTING:-1}"
DISABLE_GROUP_TEXTS="${DISABLE_GROUP_TEXTS:-1}"
SAVE_STEPS="${SAVE_STEPS:-500}"
SAVE_TOTAL_LIMIT="${SAVE_TOTAL_LIMIT:-20}"
LOGGING_STEPS="${LOGGING_STEPS:-50}"
WORKERS="${WORKERS:-4}"
PREP_WORKERS="${PREP_WORKERS:-32}"

# AR mode: MDM=0
MDM="${MDM:-0}"
BD_SIZE="${BD_SIZE:-8}"
LEARN_PADDING="${LEARN_PADDING:-0}"
CONPLE_MASK="${CONPLE_MASK:-0}"

TEMPLATE="${TEMPLATE:-qwen2_5_no_reasoning}"
RUN_ID="${RUN_ID:-vila_webds_ar_$(date +%Y%m%d_%H%M%S)}"
OUTPUT_DIR="${OUTPUT_DIR:-${PROJECT_ROOT}/output_models/${RUN_ID}}"
LOG_DIR="${LOG_DIR:-${PROJECT_ROOT}/log/${RUN_ID}}"
DS_CONFIG="${DS_CONFIG:-configs/ds_config_zero2_no_offload.json}"

# Detect number of GPUs available
NGPUS="${NGPUS:-$(python3 -c 'import torch; print(torch.cuda.device_count())' 2>/dev/null || echo 1)}"

# AutoResume
AUTO_RESUME="${AUTO_RESUME:-1}"
RESUME_FROM_CHECKPOINT="${RESUME_FROM_CHECKPOINT:-}"

echo "Debug: Current PWD=$(pwd)"
echo "Debug: OUTPUT_DIR=${OUTPUT_DIR}"
echo "Debug: LOG_DIR=${LOG_DIR}"
mkdir -p "${OUTPUT_DIR}" "${LOG_DIR}"

echo "=========================================="
echo "Configuration:"
echo "  Model: ${MODEL_PATH}"
echo "  Tokenizer: ${TOKENIZER_NAME}"
echo "  Dataset WDSMETA: ${DATASET_WDSMETA}"
echo "  Dataset Backend: ${DATASET_BACKEND}"
echo "  Output: ${OUTPUT_DIR}"
echo "  Log Dir: ${LOG_DIR}"
echo "  Epochs: ${EPOCHS}"
echo "  Block Size: ${BLOCK_SIZE}"
echo "  Batch: ${BATCH}"
echo "  Grad Acc: ${GRAD_ACC}"
echo "  LR: ${LR}"
echo "  Warmup Ratio: ${WARMUP_RATIO}"
echo "  BF16: ${BF16}"
echo "  MDM: ${MDM} (AR mode)"
echo "  DS Config: ${DS_CONFIG}"
echo "  Processor MinPixels/MaxPixels: ${MIN_PIXELS}/${MAX_PIXELS}"
echo "  Num GPUs: ${NGPUS}"
echo "  Auto Resume: ${AUTO_RESUME}"
echo "=========================================="
echo ""

# Check paths
echo "Checking paths..."
if [ ! -d "${MODEL_PATH}" ]; then
    echo "❌ Model path does not exist: ${MODEL_PATH}"
    exit 1
fi
echo "✅ Model path: ${MODEL_PATH}"

# Check if DATASET_WDSMETA is a file or a mixture name
if [ -f "${DATASET_WDSMETA}" ]; then
    echo "✅ Dataset wdsmeta.json: ${DATASET_WDSMETA}"
else
    echo "✅ Dataset mixture name: ${DATASET_WDSMETA}"
    echo "ℹ️  Will use VILA's mixture registry to load datasets"
fi

echo ""

# AutoResume: detect latest checkpoint
if [ "${AUTO_RESUME}" = "1" ] && [ -z "${RESUME_FROM_CHECKPOINT}" ]; then
    if [ -d "${OUTPUT_DIR}" ]; then
        LATEST_CHECKPOINT=$(find "${OUTPUT_DIR}" -maxdepth 1 -type d -name "checkpoint-*" | sort -V | tail -n 1)
        if [ -n "${LATEST_CHECKPOINT}" ]; then
            RESUME_FROM_CHECKPOINT="${LATEST_CHECKPOINT}"
            echo "🔄 AutoResume: Found existing checkpoint: ${RESUME_FROM_CHECKPOINT}"
        else
            echo "ℹ️ AutoResume: No checkpoint found, starting from scratch"
        fi
    else
        echo "ℹ️ AutoResume: Output directory does not exist, starting from scratch"
    fi
elif [ -n "${RESUME_FROM_CHECKPOINT}" ]; then
    echo "📍 AutoResume: Using explicitly specified checkpoint: ${RESUME_FROM_CHECKPOINT}"
fi

# Build deepspeed command
CMD=(
    deepspeed --num_gpus "${NGPUS}" examples/finetune_qwenvl.py
    --model_name_or_path "${MODEL_PATH}"
    --tokenizer_name "${TOKENIZER_NAME}"
    --conversation_template "${TEMPLATE}"
    --trust_remote_code 1
    --dataset_path "${DATASET_WDSMETA}"
    --output_dir "${OUTPUT_DIR}"
    --num_train_epochs "${EPOCHS}"
    --lr_scheduler_type constant_with_warmup
    --warmup_ratio "${WARMUP_RATIO}"
    --learning_rate "${LR}"
    --disable_group_texts "${DISABLE_GROUP_TEXTS}"
    --block_size "${BLOCK_SIZE}"
    --per_device_train_batch_size "${BATCH}"
    --gradient_accumulation_steps "${GRAD_ACC}"
    --run_name "${RUN_ID}"
    --logging_steps "${LOGGING_STEPS}"
    --do_train
    --validation_split_percentage 0
    --ddp_timeout 72000
    --save_steps "${SAVE_STEPS}"
    --save_total_limit "${SAVE_TOTAL_LIMIT}"
    --dataloader_num_workers "${WORKERS}"
    --preprocessing_num_workers "${PREP_WORKERS}"
    --bd_size "${BD_SIZE}"
    --mdm "${MDM}"
    --learn_padding "${LEARN_PADDING}"
    --conplemenrary_mask "${CONPLE_MASK}"
    --deepspeed "${DS_CONFIG}"
)

# Add resume_from_checkpoint if available
if [ -n "${RESUME_FROM_CHECKPOINT}" ]; then
    CMD+=( --resume_from_checkpoint "${RESUME_FROM_CHECKPOINT}" )
    echo "📍 Resuming from checkpoint: ${RESUME_FROM_CHECKPOINT}"
else
    CMD+=( --overwrite_output_dir )
    echo "🆕 Starting fresh training (will overwrite output dir)"
fi

if [ "${BF16}" = "1" ]; then
    CMD+=( --bf16 )
fi
if [ "${GRADIENT_CHECKPOINTING}" = "1" ]; then
    CMD+=( --gradient_checkpointing )
fi

echo ""
echo "=========================================="
echo "Starting Deepspeed training (AR mode)..."
echo "=========================================="
echo "Command: ${CMD[*]}"
echo ""

# Run training
"${CMD[@]}" 2>&1 | tee "${LOG_DIR}/train.log"
EXIT_CODE=${PIPESTATUS[0]}

echo ""
echo "=========================================="
if [ $EXIT_CODE -eq 0 ]; then
    echo "✅ VILA WebDataset AR VLM Training Success!"
    # ============================================
    # AUTO EVAL: 训练成功后自动提交评测任务
    # ============================================
    echo ""
    echo "=========================================="
    echo "🚀 Auto-submitting evaluation jobs..."
    echo "=========================================="
    
    if [ -f "${PROJECT_ROOT}/chengyue_script/dvlm/evaluation/submit_batch_eval.sh" ]; then
        bash "${PROJECT_ROOT}/chengyue_script/dvlm/evaluation/submit_batch_eval.sh" \
            "${OUTPUT_DIR}" \
            --skip-existing 1 \
            2>&1 | tee "${LOG_DIR}/auto_eval_submit.log"
        echo "✅ Evaluation jobs submitted. Check: ${OUTPUT_DIR}/eval_vlmeval/"
    else
        echo "⚠️  submit_batch_eval.sh not found, skipping auto eval"
    fi

else
    echo "❌ VILA WebDataset AR VLM Training Failed (exit code: ${EXIT_CODE})"
    echo -e "Please check logs:\n  Output: ${LOG_DIR}/train.log"
fi
echo "=========================================="

exit $EXIT_CODE

