#!/bin/bash
# ==============================================================================
# 批量提交 VLMEvalKit 评测任务
# 为指定实验目录下的所有 checkpoint 生成并提交 sbatch job
# 支持多配置组合 (USE_MDM, USE_SPECULATIVE)
#
# Usage:
#   bash chengyue_script/dvlm/evaluation/submit_batch_eval.sh <experiment_dir> [options]
#
# Examples:
#   # 评测所有checkpoint (默认配置)
#   bash chengyue_script/dvlm/evaluation/submit_batch_eval.sh output_models/xxx
#
#   # 使用所有配置组合
#   bash chengyue_script/dvlm/evaluation/submit_batch_eval.sh output_models/xxx --configs all
#
#   # 指定配置组合
#   bash chengyue_script/dvlm/evaluation/submit_batch_eval.sh output_models/xxx \
#       --configs "MDM=1_SPEC=0 MDM=1_SPEC=1"
#
#   # 只评测最新的5个checkpoint
#   bash chengyue_script/dvlm/evaluation/submit_batch_eval.sh output_models/xxx --last 5
#
#   # dry-run模式（只生成脚本不提交）
#   bash chengyue_script/dvlm/evaluation/submit_batch_eval.sh output_models/xxx --dry-run
# ==============================================================================

set -eo pipefail

# 项目根目录 (自动检测)
SCRIPT_DIR_REAL="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(cd "${SCRIPT_DIR_REAL}/../../.." && pwd)"
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

# 默认配置
CONDA_ENV="dvlm"
DEFAULT_TASKS="SEEDBench2_Plus AI2D_TEST RealWorldQA DocVQA_VAL ChartQA_TEST MathVista_MINI MMMU_DEV_VAL MMBench_DEV_EN_V11 GQA_TestDev_Balanced TextVQA_VAL POPE"

# 有效的配置组合
# MDM=0_SPEC=0: AR mode (standard autoregressive)
# MDM=1_SPEC=0: MDM mode, standard decoding
# MDM=1_SPEC=1: MDM mode, speculative decoding
# Note: MDM=0_SPEC=1 is invalid (speculative is MDM-specific)
ALL_CONFIGS="MDM=0_SPEC=0 MDM=1_SPEC=0 MDM=1_SPEC=1"
DEFAULT_CONFIGS="MDM=1_SPEC=0"  # 默认只跑 MDM 标准解码

# 解析参数
EXPERIMENT_DIR=""
LAST_N=""
CHECKPOINTS=""
DRY_RUN=0
TASKS="${TASKS:-${DEFAULT_TASKS}}"
CONFIGS="${CONFIGS:-}"
USE_MDM="${USE_MDM:-}"
USE_SPECULATIVE="${USE_SPECULATIVE:-}"
SKIP_EXISTING="${SKIP_EXISTING:-1}"
FORCE_LEFT_TO_RIGHT="${FORCE_LEFT_TO_RIGHT:-0}"

show_help() {
    echo "批量提交 VLMEvalKit 评测任务 (支持多配置组合)"
    echo ""
    echo "Usage: $0 <experiment_dir> [options]"
    echo ""
    echo "Options:"
    echo "  --configs        配置组合列表 (空格分隔) 或 'all'"
    echo "                   有效组合: MDM=0_SPEC=0, MDM=1_SPEC=0, MDM=1_SPEC=1"
    echo "                   'all' 会展开为所有有效组合"
    echo "                   默认: MDM=1_SPEC=0"
    echo "  --last N         只评测最后 N 个 checkpoint"
    echo "  --checkpoints    指定要评测的 checkpoint 列表（空格分隔）"
    echo "  --tasks          评测任务列表（默认: 全部11个benchmark）"
    echo "  --use-mdm        [向后兼容] 是否使用 MDM 模式 (会被 --configs 覆盖)"
    echo "  --use-speculative [向后兼容] 是否使用 speculative 解码 (会被 --configs 覆盖)"
    echo "  --skip-existing  跳过已有评测结果的 checkpoint (default: 1)"
    echo "  --force-left-to-right 强制使用左到右解码 (default: 0)"
    echo "  --dry-run        只生成脚本不提交"
    echo "  -h, --help       显示帮助"
    echo ""
    echo "Directory Structure:"
    echo "  output_models/exp_name/eval_vlmeval/MDM=X_SPEC=Y/checkpoint-XXX/"
    echo ""
    echo "Examples:"
    echo "  # 使用所有配置组合"
    echo "  $0 output_models/xxx --configs all"
    echo ""
    echo "  # 只测试 MDM 模式的两种解码方式"
    echo "  $0 output_models/xxx --configs \"MDM=1_SPEC=0 MDM=1_SPEC=1\""
    echo ""
    echo "  # 向后兼容: 单配置模式"
    echo "  $0 output_models/xxx --use-mdm 1 --use-speculative 0"
    exit 0
}

# 解析参数
while [[ $# -gt 0 ]]; do
    case $1 in
        --configs)
            CONFIGS="$2"
            shift 2
            ;;
        --last)
            LAST_N="$2"
            shift 2
            ;;
        --checkpoints)
            CHECKPOINTS="$2"
            shift 2
            ;;
        --tasks)
            TASKS="$2"
            shift 2
            ;;
        --use-mdm)
            USE_MDM="$2"
            shift 2
            ;;
        --use-speculative)
            USE_SPECULATIVE="$2"
            shift 2
            ;;
        --skip-existing)
            SKIP_EXISTING="$2"
            shift 2
            ;;
        --force-left-to-right)
            FORCE_LEFT_TO_RIGHT="$2"
            shift 2
            ;;
        --dry-run)
            DRY_RUN=1
            shift
            ;;
        -h|--help)
            show_help
            ;;
        *)
            if [ -z "$EXPERIMENT_DIR" ]; then
                EXPERIMENT_DIR="$1"
            else
                echo "Unknown option: $1"
                exit 1
            fi
            shift
            ;;
    esac
done

if [ -z "$EXPERIMENT_DIR" ]; then
    echo "Error: 请指定实验目录"
    show_help
fi

# 转换为绝对路径
if [[ "$EXPERIMENT_DIR" != /* ]]; then
    EXPERIMENT_DIR="${PROJECT_ROOT}/${EXPERIMENT_DIR}"
fi

if [ ! -d "$EXPERIMENT_DIR" ]; then
    echo "Error: 目录不存在: $EXPERIMENT_DIR"
    exit 1
fi

# 确定配置列表
if [ -n "$CONFIGS" ]; then
    if [ "$CONFIGS" = "all" ]; then
        CONFIG_LIST=($ALL_CONFIGS)
    else
        CONFIG_LIST=($CONFIGS)
    fi
elif [ -n "$USE_MDM" ] || [ -n "$USE_SPECULATIVE" ]; then
    # 向后兼容: 从 USE_MDM 和 USE_SPECULATIVE 构建配置
    MDM_VAL="${USE_MDM:-1}"
    SPEC_VAL="${USE_SPECULATIVE:-0}"
    CONFIG_LIST=("MDM=${MDM_VAL}_SPEC=${SPEC_VAL}")
else
    # 使用默认配置
    CONFIG_LIST=($DEFAULT_CONFIGS)
fi

# 验证配置有效性
validate_config() {
    local cfg="$1"
    local mdm spec
    mdm=$(echo "$cfg" | grep -oP 'MDM=\K[01]' || echo "")
    spec=$(echo "$cfg" | grep -oP 'SPEC=\K[01]' || echo "")
    
    if [ -z "$mdm" ] || [ -z "$spec" ]; then
        echo "Error: 无效配置格式: $cfg (应为 MDM=X_SPEC=Y)"
        return 1
    fi
    
    # MDM=0_SPEC=1 无效
    if [ "$mdm" = "0" ] && [ "$spec" = "1" ]; then
        echo "Error: 无效配置: $cfg (MDM=0 时不支持 speculative 解码)"
        return 1
    fi
    
    return 0
}

for cfg in "${CONFIG_LIST[@]}"; do
    if ! validate_config "$cfg"; then
        exit 1
    fi
done

EXPERIMENT_NAME=$(basename "$EXPERIMENT_DIR")
echo "=============================================="
echo "批量提交评测任务 (多配置模式)"
echo "=============================================="
echo "实验目录: $EXPERIMENT_DIR"
echo "实验名称: $EXPERIMENT_NAME"
echo "配置列表: ${CONFIG_LIST[*]}"
echo "TASKS: $TASKS"
echo "SKIP_EXISTING: $SKIP_EXISTING"
echo "FORCE_LEFT_TO_RIGHT: $FORCE_LEFT_TO_RIGHT"
echo "DRY_RUN: $DRY_RUN"
echo ""

# 获取 checkpoint 列表
if [ -n "$CHECKPOINTS" ]; then
    # 使用指定的 checkpoint
    CKPT_LIST=($CHECKPOINTS)
else
    # 自动扫描 checkpoint
    CKPT_LIST=($(find "$EXPERIMENT_DIR" -maxdepth 1 -type d -name "checkpoint-*" | sort -V))
    
    if [ -n "$LAST_N" ]; then
        # 只取最后 N 个
        CKPT_LIST=("${CKPT_LIST[@]: -$LAST_N}")
    fi
fi

if [ ${#CKPT_LIST[@]} -eq 0 ]; then
    echo "Error: 没有找到 checkpoint"
    exit 1
fi

echo "找到 ${#CKPT_LIST[@]} 个 checkpoint:"
for ckpt in "${CKPT_LIST[@]}"; do
    echo "  - $(basename $ckpt)"
done
echo ""

# 创建 eval job 目录
EVAL_DIR="${EXPERIMENT_DIR}/eval_vlmeval"
mkdir -p "$EVAL_DIR"

# 统计
SUBMITTED=0
SKIPPED=0
set +e  # 允许计数操作失败不退出

# 为每个配置和每个 checkpoint 生成并提交 job
for CONFIG in "${CONFIG_LIST[@]}"; do
    echo "=============================================="
    echo "配置: ${CONFIG}"
    echo "=============================================="
    
    # 解析配置
    CFG_USE_MDM=$(echo "$CONFIG" | grep -oP 'MDM=\K[01]')
    CFG_USE_SPECULATIVE=$(echo "$CONFIG" | grep -oP 'SPEC=\K[01]')
    
    # 创建配置目录
    CONFIG_DIR="${EVAL_DIR}/${CONFIG}"
    mkdir -p "$CONFIG_DIR"
    
    for CKPT_PATH in "${CKPT_LIST[@]}"; do
        CKPT_NAME=$(basename "$CKPT_PATH")
        OUTPUT_DIR="${CONFIG_DIR}/${CKPT_NAME}"
        JOB_NAME="${EXPERIMENT_NAME}_${CONFIG}_${CKPT_NAME}"
        # 截断 job name 避免超过 SLURM 限制
        JOB_NAME="${JOB_NAME:0:63}"
        SCRIPT_PATH="${OUTPUT_DIR}/run_eval.sh"
        
        # 检查是否已有评测结果
        if [ "$SKIP_EXISTING" = "1" ] && [ -f "${OUTPUT_DIR}/summary.md" ]; then
            echo "⏭️  跳过 ${CONFIG}/${CKPT_NAME} (已有评测结果)"
            ((SKIPPED++))
            continue
        fi
        
        # 创建输出目录
        mkdir -p "$OUTPUT_DIR"
        
        # 生成 sbatch 脚本
        cat > "$SCRIPT_PATH" << EOF
#!/bin/bash
#SBATCH -A nvr_lpr_llm
#SBATCH -p batch_block1,batch_singlenode
#SBATCH -t 4:00:00
#SBATCH -N 1
#SBATCH -J eval_${JOB_NAME}
#SBATCH --gpus-per-node 8
#SBATCH --cpus-per-task 16
#SBATCH --mem=512G
#SBATCH -o ${OUTPUT_DIR}/slurm_out.log
#SBATCH -e ${OUTPUT_DIR}/slurm_err.log

cd ${PROJECT_ROOT}
source "${PROJECT_ROOT}/scripts/cluster_env.sh"
lmflow_activate_conda

echo "=============================================="
echo "Evaluating: ${CKPT_NAME}"
echo "Config: ${CONFIG}"
echo "  USE_MDM: ${CFG_USE_MDM}"
echo "  USE_SPECULATIVE: ${CFG_USE_SPECULATIVE}"
echo "Model: ${CKPT_PATH}"
echo "Output: ${OUTPUT_DIR}"
echo "=============================================="

# Run VLMEvalKit evaluation
MODEL_PATH="${CKPT_PATH}" \\
NUM_GPUS=8 \\
TASKS="${TASKS}" \\
OUTPUT_BASE="${OUTPUT_DIR}" \\
USE_MDM=${CFG_USE_MDM} \\
USE_SPECULATIVE=${CFG_USE_SPECULATIVE} \\
bash chengyue_script/dvlm/evaluation/vlmeval/run_all_benchmarks.sh

# Mark evaluation as complete
touch "${OUTPUT_DIR}/.eval_done"

echo "=============================================="
echo "Evaluation complete: ${CONFIG}/${CKPT_NAME}"
echo "=============================================="
EOF
        
        chmod +x "$SCRIPT_PATH"
        
        if [ "$DRY_RUN" = "1" ]; then
            echo "📝 生成脚本: $SCRIPT_PATH"
        else
            # 提交 job
            JOB_ID=$(sbatch "$SCRIPT_PATH" | awk '{print $NF}')
            echo "✅ 提交 ${CONFIG}/${CKPT_NAME} -> Job ID: ${JOB_ID}"
        fi
        ((SUBMITTED++))
    done
done

echo ""
echo "=============================================="
echo "汇总"
echo "=============================================="
echo "配置数: ${#CONFIG_LIST[@]}"
echo "Checkpoint 数: ${#CKPT_LIST[@]}"
echo "提交/生成: ${SUBMITTED}"
echo "跳过: ${SKIPPED}"
echo "输出目录: ${EVAL_DIR}"
echo ""
echo "目录结构:"
for cfg in "${CONFIG_LIST[@]}"; do
    echo "  ${EVAL_DIR}/${cfg}/checkpoint-*/"
done

if [ "$DRY_RUN" = "1" ]; then
    echo ""
    echo "提示: 使用以下命令批量提交:"
    echo "  for f in ${EVAL_DIR}/*/checkpoint-*/run_eval.sh; do sbatch \$f; done"
fi

echo "=============================================="
