#!/usr/bin/env bash
set -euo pipefail

CANN_ENV=${CANN_ENV:-/usr/local/Ascend/cann-8.5.1/cann-8.5.1/set_env.sh}
VENV_ACTIVATE=${VENV_ACTIVATE:-/data/vllm-ascend-v0.18.0/.venv/bin/activate}
source "${CANN_ENV}"
source "${VENV_ACTIVATE}"

SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
REPO_ROOT="$(cd -- "${SCRIPT_DIR}/.." && pwd)"

VLLM_ASCEND_CUSTOM_OP_ROOT=${VLLM_ASCEND_CUSTOM_OP_ROOT-/data/vllm-ascend-v0.18.0/vllm_ascend/_cann_ops_custom/vendors/vllm-ascend}
if [[ -n "${VLLM_ASCEND_CUSTOM_OP_ROOT}" ]]; then
    export ASCEND_CUSTOM_OPP_PATH="${VLLM_ASCEND_CUSTOM_OP_ROOT}${ASCEND_CUSTOM_OPP_PATH:+:${ASCEND_CUSTOM_OPP_PATH}}"
    export LD_LIBRARY_PATH="${VLLM_ASCEND_CUSTOM_OP_ROOT}/op_api/lib${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}"
fi
export CANN_ENV VENV_ACTIVATE VLLM_ASCEND_CUSTOM_OP_ROOT

export VLLM_TARGET_DEVICE=ascend
export VLLM_USE_MODELSCOPE=True
export ASCEND_RT_VISIBLE_DEVICES=${ASCEND_RT_VISIBLE_DEVICES:-0}
export NPU_INDEX=${NPU_INDEX:-$(python "${SCRIPT_DIR}/detect_npu_index.py")}
export NPU_CHIP_INDEX=${NPU_CHIP_INDEX:-0}
export MODELSCOPE_CACHE=${MODELSCOPE_CACHE:-/data/modelscope-cache}
export HF_HOME=${HF_HOME:-/data/huggingface-cache}
export TRITON_CACHE_DIR=${TRITON_CACHE_DIR:-/data/triton-cache-cann851}

MODEL_PATH=${MODEL_PATH:?MODEL_PATH must point to the model}
DATA_DIR=${DATA_DIR:-/data/accuracy-gsm8k-gpqa}
RESULTS_DIR=${RESULTS_DIR:-/data/accuracy-results}
ACCURACY_LIMIT=${ACCURACY_LIMIT:-0}
SEED=${SEED:-0}
DATASET_REVISION=${DATASET_REVISION:-master}
MAX_MODEL_LEN=${MAX_MODEL_LEN:-4096}
GPU_MEMORY_UTILIZATION=${GPU_MEMORY_UTILIZATION:-0.5}
GSM8K_MAX_TOKENS=${GSM8K_MAX_TOKENS:-128}
GPQA_MAX_TOKENS=${GPQA_MAX_TOKENS:-16}
MAX_ACCURACY_DROP=${MAX_ACCURACY_DROP:-0}

if [[ ${SKIP_NPU_HEALTH_CHECK:-0} != 1 ]]; then
    bash "${SCRIPT_DIR}/check_npu_health.sh"
fi

run_case() {
    local case_name=$1
    local group_size=$2
    local num_in_group=$3
    local case_dir="${RESULTS_DIR}/${case_name}"

    python "${REPO_ROOT}/scripts/accuracy_gsm8k_gpqa.py" \
        --model "$MODEL_PATH" --dataset gsm8k --data-dir "$DATA_DIR" --output-dir "$case_dir" \
        --limit "$ACCURACY_LIMIT" --seed "$SEED" --dataset-revision "$DATASET_REVISION" \
        --max-model-len "$MAX_MODEL_LEN" \
        --gpu-memory-utilization "$GPU_MEMORY_UTILIZATION" \
        --max-tokens "$GSM8K_MAX_TOKENS" --offload-group-size "$group_size" \
        --offload-num-in-group "$num_in_group"
    python "${REPO_ROOT}/scripts/accuracy_gsm8k_gpqa.py" \
        --model "$MODEL_PATH" --dataset gpqa --data-dir "$DATA_DIR" --output-dir "$case_dir" \
        --limit "$ACCURACY_LIMIT" --seed "$SEED" --dataset-revision "$DATASET_REVISION" \
        --max-model-len "$MAX_MODEL_LEN" \
        --gpu-memory-utilization "$GPU_MEMORY_UTILIZATION" \
        --max-tokens "$GPQA_MAX_TOKENS" --offload-group-size "$group_size" \
        --offload-num-in-group "$num_in_group"
}

run_case "baseline_0_1" 0 1
run_case "group_8_num_1" 8 1
run_case "group_8_num_2" 8 2

python "${REPO_ROOT}/scripts/verify_accuracy_matrix.py" \
    --results-dir "$RESULTS_DIR" \
    --max-accuracy-drop "$MAX_ACCURACY_DROP"
