#!/bin/bash

# Project root
PROJECT_ROOT="/lustre/fs11/portfolios/nvr/projects/nvr_elm_llm/users/chengyuew/workspace/shiyi/LMFlow"

# 训练参数设置 (from ar_baseline)
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}"

DATASET_JSON="${DATASET_JSON:-/lustre/fsw/portfolios/nvr/projects/nvr_elm_llm/dataset/vila-sft/ShareGPT4V/jason-filter-sharegpt4v_mix665k_cap23k_coco-ap9k_lcs3k_sam9k_div2k.json}"
IMAGE_FOLDER="${IMAGE_FOLDER:-/lustre/fsw/portfolios/nvr/projects/nvr_elm_llm/dataset/vila-sft/ShareGPT4V/data}"

EPOCHS="${EPOCHS:-1}"
BLOCK_SIZE="${BLOCK_SIZE:-2048}"
BATCH="${BATCH:-2}"
GRAD_ACC="${GRAD_ACC:-4}"
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:-100}"
SAVE_TOTAL_LIMIT="${SAVE_TOTAL_LIMIT:-20}"
LOGGING_STEPS="${LOGGING_STEPS:-50}"
WORKERS="${WORKERS:-8}"
PREP_WORKERS="${PREP_WORKERS:-32}"
MDM="${MDM:-0}"
BD_SIZE="${BD_SIZE:-8}"
LEARN_PADDING="${LEARN_PADDING:-1}"
CONPLE_MASK="${CONPLE_MASK:-1}"
TEMPLATE="${TEMPLATE:-qwen2_5_no_reasoning}"

RUN_ID="${RUN_ID:-sharegpt4v_ar_baseline_single_node}"
OUTPUT_DIR="${OUTPUT_DIR:-${PROJECT_ROOT}/output_models/${RUN_ID}}"
DS_CONFIG="${DS_CONFIG:-configs/ds_config_zero2_no_offload.json}"

# Qwen2.5-VL processor dynamic resolution controls
# export MIN_PIXELS="${MIN_PIXELS:-200704}"
# export MAX_PIXELS="${MAX_PIXELS:-200704}"
# export RESIZED_HEIGHT="${RESIZED_HEIGHT:-}"
# export RESIZED_WIDTH="${RESIZED_WIDTH:-}"

# Environment variables
# export CUDA_HOME=/lustre/fsw/portfolios/nvr/users/chengyuew/anaconda3/envs/lmflow

trust_remote_code=1

job_name=train_sharegpt4v_vlm_ar_baseline_single_node
exp_id=train_sharegpt4v_vlm_ar_baseline_single_node

# Print training info
echo "=========================================="
echo "MDM_Qwen2VL VLM Training (ShareGPT4V, single-node)"
echo "=========================================="
echo "Date: $(date)"
echo ""

# Auto-detect latest checkpoint (AutoResume)
latest_checkpoint=""
if [ -d "${OUTPUT_DIR}" ]; then
    # Find all checkpoint directories, sort by number, get the latest one
    latest_checkpoint=$(find "${OUTPUT_DIR}" -maxdepth 1 -name "checkpoint-*" -type d | sort -V | tail -1)
    if [ -n "${latest_checkpoint}" ]; then
        echo "🔄 Found latest checkpoint: ${latest_checkpoint}"
    else
        echo "ℹ️ No checkpoint found in ${OUTPUT_DIR}"
        latest_checkpoint=""
    fi
else
    echo "ℹ️ Output directory ${OUTPUT_DIR} does not exist, will train from scratch"
    latest_checkpoint=""
fi

# Create necessary directories
log_dir=log/${job_name}
mkdir -p ${OUTPUT_DIR} ${log_dir}

# Build resume_from_checkpoint argument
resume_arg=""
if [ -n "${latest_checkpoint}" ]; then
    resume_arg="--resume_from_checkpoint ${latest_checkpoint}"
    echo "📍 Will resume training from checkpoint"
else
    resume_arg="--overwrite_output_dir"
    echo "🆕 Starting training from scratch"
fi

echo ""
echo "=========================================="
echo "Configuration:"
echo "  Model: ${MODEL_PATH}"
echo "  Tokenizer: ${TOKENIZER_NAME}"
echo "  Dataset JSON: ${DATASET_JSON}"
echo "  Image Folder: ${IMAGE_FOLDER}"
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}  BD_SIZE: ${BD_SIZE}  LearnPadding: ${LEARN_PADDING}  ComplementaryMask: ${CONPLE_MASK}"
echo "  DS Config: ${DS_CONFIG}"
echo "  Processor MinPixels/MaxPixels: ${MIN_PIXELS}/${MAX_PIXELS}"
echo "  Template: ${TEMPLATE}"
echo "=========================================="
echo ""

# GPU monitoring (run in background)
nvidia-smi --query-gpu=timestamp,memory.used,memory.total,utilization.gpu,pstate \
           --format=csv,noheader,nounits -l 60 > ${log_dir}/gpu_monitor.log 2>&1 &
GPU_MONITOR_PID=$!

# Activate conda environment
source "${PROJECT_ROOT}/scripts/cluster_env.sh"
lmflow_activate_conda

# Change to project directory
cd ${PROJECT_ROOT}

# Run DeepSpeed training
echo "🚀 Starting training..."
echo ""

deepspeed --master_port=23456 \
    examples/finetune_qwenvl.py \
    --model_name_or_path ${MODEL_PATH} \
    --tokenizer_name ${TOKENIZER_NAME} \
    --conversation_template ${TEMPLATE} \
    --trust_remote_code ${trust_remote_code} \
    --dataset_path ${DATASET_JSON} \
    --image_folder ${IMAGE_FOLDER} \
    --output_dir ${OUTPUT_DIR} \
    ${resume_arg} \
    --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 ${exp_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} \
    $([ "${BF16}" = "1" ] && echo "--bf16") \
    $([ "${GRADIENT_CHECKPOINTING}" = "1" ] && echo "--gradient_checkpointing") \
    2>&1 | tee ${log_dir}/training.log

exit_code=${PIPESTATUS[0]}

# Stop GPU monitoring
kill ${GPU_MONITOR_PID} 2>/dev/null || true

echo ""
echo "=========================================="
if [ $exit_code -eq 0 ]; then
    echo "✅ ShareGPT4V 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 "❌ ShareGPT4V VLM Training Failed (exit code: ${exit_code})"
fi
echo "=========================================="
echo ""

exit $exit_code

