#!/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)"

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 XDG_CACHE_HOME=${XDG_CACHE_HOME:-/data/cache}
export HF_HOME=${HF_HOME:-/data/huggingface-cache}
export MODELSCOPE_CACHE=${MODELSCOPE_CACHE:-/data/modelscope-cache}
export VLLM_CACHE_ROOT=${VLLM_CACHE_ROOT:-/data/vllm-cache}
export PYTORCH_NPU_ALLOC_CONF=${PYTORCH_NPU_ALLOC_CONF:-expandable_segments:True}
export TRITON_CACHE_DIR=${TRITON_CACHE_DIR:-/data/triton-cache-cann851}

MODEL_PATH=${MODEL_PATH:?MODEL_PATH must point to the model}
RESULTS_DIR=${RESULTS_DIR:-/data/performance-results}
REPEATS=${REPEATS:-3}
WARMUP_ROUNDS=${WARMUP_ROUNDS:-1}
MEASURED_ROUNDS=${MEASURED_ROUNDS:-3}
PROMPT_COUNT=${PROMPT_COUNT:-8}
INPUT_TOKENS=${INPUT_TOKENS:-128}
MAX_MODEL_LEN=${MAX_MODEL_LEN:-512}
MAX_TOKENS=${MAX_TOKENS:-64}
GPU_MEMORY_UTILIZATION=${GPU_MEMORY_UTILIZATION:-0.5}
SEED=${SEED:-0}

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

mkdir -p "${RESULTS_DIR}"
RAW_RESULTS="${RESULTS_DIR}/performance-runs.jsonl"
RUN_LOG="${RESULTS_DIR}/performance-run.log"
: > "${RAW_RESULTS}"
: > "${RUN_LOG}"

CASES=(
    "baseline_0_1:0:1"
    "group_8_num_1:8:1"
    "group_8_num_2:8:2"
)

for ((repeat=0; repeat<REPEATS; repeat++)); do
    start=$((repeat % ${#CASES[@]}))
    for ((offset=0; offset<${#CASES[@]}; offset++)); do
        index=$(((start + offset) % ${#CASES[@]}))
        IFS=: read -r label group_size num_in_group <<< "${CASES[index]}"
        echo "repeat=${repeat} case=${label}" | tee -a "${RUN_LOG}"
        python "${SCRIPT_DIR}/bench_one.py" \
            --model "${MODEL_PATH}" \
            --label "${label}" \
            --offload-group-size "${group_size}" \
            --offload-num-in-group "${num_in_group}" \
            --prompt-count "${PROMPT_COUNT}" \
            --input-tokens "${INPUT_TOKENS}" \
            --max-tokens "${MAX_TOKENS}" \
            --max-model-len "${MAX_MODEL_LEN}" \
            --warmup-rounds "${WARMUP_ROUNDS}" \
            --measured-rounds "${MEASURED_ROUNDS}" \
            --seed "${SEED}" \
            --result-file "${RAW_RESULTS}" \
            --gpu-memory-utilization "${GPU_MEMORY_UTILIZATION}" \
            --npu-index "${NPU_INDEX}" 2>&1 | tee -a "${RUN_LOG}"
    done
done

python "${SCRIPT_DIR}/bench_param_matrix.py" \
    --input "${RAW_RESULTS}" \
    --output-dir "${RESULTS_DIR}" \
    --baseline-label baseline_0_1
