set -ex
export CUDA_VISIBLE_DEVICES="0,1,2,3"

PROMPT_TYPE=qwen-base-boxed
EXP=target_0.5-n8x4
BASE_MODEL=Qwen/Qwen2.5-Math-7B
LAST_STEP=400
max_tokens=2048
top_p=0.95

MODEL_NAME_OR_PATH=$(pwd)/checkpoints/${BASE_MODEL##*/}/$EXP/global_step_$LAST_STEP/actor
OUTPUT_DIR=results/${BASE_MODEL##*/}/$EXP
benchmarks="gsm8k,math500,minerva_math,olympiadbench,college_math,aime24,amc23,arc_c,gpqa,mmlu_pro"
SPLIT="test"
NUM_TEST_SAMPLE=-1
OVERWRITE=true

 python scripts/model_merger.py \
     --backend fsdp \
     --tie-word-embedding \
     --hf_model_path $BASE_MODEL \
     --local_dir $MODEL_NAME_OR_PATH

cd simpleRL-reason/examples/simplelr_math_eval

# English open datasets
DATA_NAME=${benchmarks}

if [ "$OVERWRITE" = "true" ]; then
    OVERWRITE="--overwrite"
else
    OVERWRITE=""
fi
# Split benchmarks into two groups
IFS=',' read -ra BENCHMARK_ARRAY <<< "$benchmarks"
REGULAR_BENCHMARKS=()
SPECIAL_BENCHMARKS=()

for benchmark in "${BENCHMARK_ARRAY[@]}"; do
    if [[ "$benchmark" == "aime24" || "$benchmark" == "amc23" ]]; then
        SPECIAL_BENCHMARKS+=("$benchmark")
    else
        REGULAR_BENCHMARKS+=("$benchmark")
    fi
done


# Run regular benchmarks with n_sampling=1
if [ ${#REGULAR_BENCHMARKS[@]} -gt 0 ]; then
    REGULAR_BENCHMARKS_STR=$(IFS=,; echo "${REGULAR_BENCHMARKS[*]}")
    TOKENIZERS_PARALLELISM=false \
    python -u math_eval.py \
        --model_name_or_path ${MODEL_NAME_OR_PATH} \
        --data_name ${REGULAR_BENCHMARKS_STR} \
        --output_dir ${OUTPUT_DIR} \
        --split ${SPLIT} \
        --prompt_type ${PROMPT_TYPE} \
        --num_test_sample ${NUM_TEST_SAMPLE} \
        --max_tokens_per_call ${max_tokens} \
        --seed 0 \
        --temperature 0. \
        --n_sampling 1 \
        --top_p ${top_p} \
        --start 0 \
        --end -1 \
        --use_vllm \
        --save_outputs \
        ${OVERWRITE}
fi

# Run special benchmarks (aime24, amc23) with n_sampling=32 temp=0.6
if [ ${#SPECIAL_BENCHMARKS[@]} -gt 0 ]; then
    SPECIAL_BENCHMARKS_STR=$(IFS=,; echo "${SPECIAL_BENCHMARKS[*]}")
    TOKENIZERS_PARALLELISM=false \
    python -u math_eval.py \
        --model_name_or_path ${MODEL_NAME_OR_PATH} \
        --data_name ${SPECIAL_BENCHMARKS_STR} \
        --output_dir ${OUTPUT_DIR} \
        --split ${SPLIT} \
        --prompt_type ${PROMPT_TYPE} \
        --num_test_sample ${NUM_TEST_SAMPLE} \
        --max_tokens_per_call ${max_tokens} \
        --seed 0 \
        --temperature 0.6 \
        --n_sampling 32 \
        --top_p ${top_p} \
        --start 0 \
        --end -1 \
        --use_vllm \
        --save_outputs \
        ${OVERWRITE}
fi
