#!/bin/bash
# ============================================
# Interactive Training Script for MDM_Qwen2.5-VL
# Run directly from command line (no SLURM)
# ============================================

# Signal handler for graceful shutdown (Ctrl+C)
TRAIN_PID=""

handle_interrupt() {
    echo ""
    echo "=========================================="
    echo "Received interrupt signal (Ctrl+C)"
    echo "Gracefully stopping training for checkpoint save..."
    echo "=========================================="
    
    # Send SIGTERM to training process to trigger graceful shutdown
    if [ -n "$TRAIN_PID" ] && kill -0 $TRAIN_PID 2>/dev/null; then
        echo "Sending SIGTERM to training process (PID: $TRAIN_PID)..."
        kill -TERM $TRAIN_PID 2>/dev/null || true
        # Wait for process to finish saving checkpoint
        wait $TRAIN_PID 2>/dev/null || true
    fi
    
    echo "Training interrupted."
    exit 0
}

# Set up signal handler for Ctrl+C
trap handle_interrupt INT TERM

# Resolve project root and change to it immediately
PROJECT_ROOT="$PWD"
cd "${PROJECT_ROOT}"

echo "=========================================="
echo "MDM_Qwen2.5-VL Training (MMFineReason Mixture)"
echo "=========================================="
echo "Mode: Interactive"
echo "Node: $(hostname)"
echo "Date: $(date)"
echo "Working Directory: $(pwd)"
echo ""

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

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 "nvidia-smi:"
nvidia-smi || true
echo ""

# Set environment variables
export PYTHONPATH="${PROJECT_ROOT}/src:${PROJECT_ROOT}/third_party/vila:${PROJECT_ROOT}/examples:${PYTHONPATH:-}"
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 MDM_VL_KEEP_VISION=1
export TOKENIZERS_PARALLELISM=false

# VILA dataset registry (draco-oci-iad has the correct paths for this cluster)
# VILA_DATASETS — loaded from cluster_env.sh
# DATASET_BACKEND — loaded from cluster_env.sh
# Qwen2.5-VL processor dynamic resolution controls
# Default to ~448x448 (200704 pixels) to stay aligned with 14x14 patching.
# Override via env vars only after profiling memory usage.
# export MIN_PIXELS="${MIN_PIXELS:-200704}"
# export MAX_PIXELS="${MAX_PIXELS:-200704}"

# Configurable parameters (override via env before running)
# NOTE: Must use MDM version of the model!
MODEL_PATH="${MODEL_PATH:-${PROJECT_ROOT}/src/models/Qwen2.5-VL-3B-Instruct-mdm}"
TOKENIZER_NAME="${TOKENIZER_NAME:-Qwen/Qwen2.5-VL-3B-Instruct}"

# VILA mixture name or HuggingFace dataset path
DATASET_MIXTURE="${DATASET_MIXTURE:-MMFineReason/SFT-586K}"

EPOCHS="${EPOCHS:-1}"
BATCH="${BATCH:-1}"
GRAD_ACC="${GRAD_ACC:-1}"
LR="${LR:-2e-5}"
SAVE_STEPS="${SAVE_STEPS:-500}"
SAVE_TOTAL_LIMIT="${SAVE_TOTAL_LIMIT:-20}"
LOGGING_STEPS="${LOGGING_STEPS:-1}"
WORKERS="${WORKERS:-4}"
MAX_LENGTH="${MAX_LENGTH:-2048}"
# MDM specific parameters
MDM="${MDM:-1}"
BD_SIZE="${BD_SIZE:-32}"
BLOCK_SIZE="${BLOCK_SIZE:-8192}"
LEARN_PADDING="${LEARN_PADDING:-1}"
CONPLE_MASK="${CONPLE_MASK:-1}"
ALWAYS_MASK_IM_END="${ALWAYS_MASK_IM_END:-1}"
# Enable Activation Checkpointing (saves memory by recomputing activations)
ACTIVATION_CHECKPOINTING="${ACTIVATION_CHECKPOINTING:-1}"

TEMPLATE="${TEMPLATE:-qwen2_5_no_reasoning}"
# Use fixed RUN_ID for auto-resume across runs
# Override with RUN_ID env var if you want a different experiment name
RUN_ID="${RUN_ID:-mmfinereason_mdm_qwen_init}"
OUTPUT_DIR="${OUTPUT_DIR:-${PROJECT_ROOT}/output_models/${RUN_ID}}"
LOG_DIR="${LOG_DIR:-${PROJECT_ROOT}/log/${RUN_ID}}"
NGPUS="${NGPUS:-8}"
USE_BLOCK_CAUSAL_MASK="${USE_BLOCK_CAUSAL_MASK:-1}"

# AutoResume: check for existing checkpoints
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}"

# Sanity-check pixel settings stay compatible with the model patch size (14x14).
if [[ "${MIN_PIXELS}" =~ ^[0-9]+$ ]] && (( MIN_PIXELS % 196 != 0 )); then
    echo "Warning: MIN_PIXELS=${MIN_PIXELS} is not divisible by 14*14; expect padding overhead."
fi
if [[ "${MAX_PIXELS}" =~ ^[0-9]+$ ]] && (( MAX_PIXELS % 196 != 0 )); then
    echo "Warning: MAX_PIXELS=${MAX_PIXELS} is not divisible by 14*14; expect padding overhead."
fi

echo "=========================================="
echo "Configuration (MDM Mode):"
echo "  Model: ${MODEL_PATH}"
echo "  Tokenizer: ${TOKENIZER_NAME}"
echo "  Dataset Mixture: ${DATASET_MIXTURE}"
echo "  VILA Registry: ${VILA_DATASETS}"
echo "  Output: ${OUTPUT_DIR}"
echo "  Log Dir: ${LOG_DIR}"
echo "  Epochs: ${EPOCHS}"
echo "  Batch: ${BATCH}"
echo "  Grad Acc: ${GRAD_ACC}"
echo "  LR: ${LR}"
echo "  MDM: ${MDM}  BD_SIZE: ${BD_SIZE}  LearnPadding: ${LEARN_PADDING}  ComplementaryMask: ${CONPLE_MASK}"
echo "  Processor MinPixels/MaxPixels: ${MIN_PIXELS}/${MAX_PIXELS}"
echo "  FSDP: shard_grad_op (ZeRO-2 equivalent)"
echo "  Num GPUs: ${NGPUS}"
echo "  Auto Resume: ${AUTO_RESUME}"
echo "  Use Block Causal Mask: ${USE_BLOCK_CAUSAL_MASK}"
echo "=========================================="
echo ""

# Check model path
echo "Checking paths..."
if [ ! -d "${MODEL_PATH}" ]; then
    echo "❌ Model path does not exist: ${MODEL_PATH}"
    echo "⚠️  Make sure you're using the MDM version of the model!"
    exit 1
fi
echo "✅ Model path: ${MODEL_PATH}"
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
fi

# Activation Checkpointing configuration
ac_arg=""
fsdp_wrap_arg=""
if [ "${ACTIVATION_CHECKPOINTING}" = "1" ]; then
    # In HF Trainer, --gradient_checkpointing enables activation checkpointing
    ac_arg="--gradient_checkpointing"
    # NOTE: We rely on the model's _no_split_modules attribute for auto-wrapping
    # Explicitly specifying the class name via CLI fails because the class isn't in the global namespace
    fsdp_wrap_arg="" 
    echo "✅ Activation Checkpointing enabled (relying on model._no_split_modules for FSDP wrapping)"
fi

# Build torchrun command with FSDP (as per working command)
CMD=(
    torchrun --nproc_per_node="${NGPUS}" --master_port=29503
    examples/finetune_qwenvl.py
    --model_name_or_path "${MODEL_PATH}"
    --tokenizer_name "${TOKENIZER_NAME}"
    --conversation_template "${TEMPLATE}"
    --trust_remote_code 1
    --dataset_path "${DATASET_MIXTURE}"
    --output_dir "${OUTPUT_DIR}"
    --num_train_epochs "${EPOCHS}"
    --learning_rate "${LR}"
    --per_device_train_batch_size "${BATCH}"
    --gradient_accumulation_steps "${GRAD_ACC}"
    --logging_steps "${LOGGING_STEPS}"
    --save_steps "${SAVE_STEPS}"
    --save_total_limit "${SAVE_TOTAL_LIMIT}"
    --do_train
    --dataloader_num_workers "${WORKERS}"
    --ddp_timeout 7200
    --mdm "${MDM}"
    --bd_size "${BD_SIZE}"
    --block_size "${BLOCK_SIZE}"
    --learn_padding "${LEARN_PADDING}"
    --complementary_mask "${CONPLE_MASK}"
    --use_block_causal_mask "${USE_BLOCK_CAUSAL_MASK}"
    --bf16
    --fsdp "shard_grad_op"
    --always_mask_im_end "${ALWAYS_MASK_IM_END}"
    --max_length "${MAX_LENGTH}"
    ${ac_arg}
    ${fsdp_wrap_arg}
)

# 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

echo ""
echo "=========================================="
echo "Starting MDM training with FSDP..."
echo "Command: ${CMD[*]}"
echo "=========================================="

# Run training in background to allow signal handling
"${CMD[@]}" 2>&1 | tee "${LOG_DIR}/train.log" &
TRAIN_PID=$!

# Wait for training to complete (or be interrupted by signal)
wait $TRAIN_PID
exit_code=$?

echo ""
echo "=========================================="
if [ $exit_code -eq 0 ]; then
    echo "✅ MMFineReason MDM 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 "❌ MMFineReason MDM Training Failed (exit code: ${exit_code})"
    echo -e "Please check logs:\n  Output: ${LOG_DIR}/train.log"
fi
echo "=========================================="

exit $exit_code
