#!/usr/bin/env bash
# Purpose: run a standard benchmark suite against a served model the same way every time -
#          same tasks, same shot count, same chat-template decision, repeated so the run-to-run
#          spread is measured rather than assumed - and record the settings the harness cannot
#          discover for itself, such as which quantisation is actually behind the alias
# Platform: all (spark, strix, mac, nvidia); the harness runs on the CPU and talks HTTP
# Minimum memory: 8 GB on the machine serving the model; this script needs very little
# Assumes: lm-evaluation-harness installed in a virtual environment (HARNESS), an
#          OpenAI-compatible endpoint at BASE_URL - the Part 9 gateway, llama-server or vLLM -
#          and a tokeniser identifier the harness can load for prompt length accounting
#
# Usage: bash run-suite.sh MODEL_ALIAS [TASKS]
#   MODEL_ALIAS  the name the endpoint answers to, e.g. local/chat
#   TASKS        comma-separated harness task names; defaults to ifeval
#
# Environment: BASE_URL (default http://127.0.0.1:4000/v1), API_KEY, HARNESS (venv directory),
#              BACKEND (local-chat-completions or local-completions), TOKENIZER,
#              NUM_FEWSHOT (unset: use the task file's own), CHAT_TEMPLATE (default yes),
#              RUNS (default 3),
#              LIMIT (leave unset for a full run), CONCURRENCY (default 1),
#              GEN_KWARGS (passed to --gen_kwargs; leave unset to use the task file's own),
#              QUANT (what is actually loaded), ENGINE, ENGINE_VERSION, OUT_DIR
set -euo pipefail

ALIAS="${1:-}"
TASKS="${2:-ifeval}"
BASE_URL="${BASE_URL:-http://127.0.0.1:4000/v1}"
BACKEND="${BACKEND:-local-chat-completions}"
NUM_FEWSHOT="${NUM_FEWSHOT:-}"
CHAT_TEMPLATE="${CHAT_TEMPLATE:-yes}"
RUNS="${RUNS:-3}"
CONCURRENCY="${CONCURRENCY:-1}"
QUANT="${QUANT:-unrecorded}"
ENGINE="${ENGINE:-unrecorded}"
ENGINE_VERSION="${ENGINE_VERSION:-unrecorded}"
OUT_DIR="${OUT_DIR:-suite-results}"

die() { echo "run-suite: $*" >&2; exit 1; }

[[ -n "$ALIAS" ]] || die "usage: bash run-suite.sh MODEL_ALIAS [TASKS]"

if [[ -n "${HARNESS:-}" ]]; then
  # shellcheck source=/dev/null
  source "$HARNESS/bin/activate"
fi
command -v lm_eval >/dev/null || die "lm_eval is not on PATH; set HARNESS to the virtual environment holding it"

if [[ "$QUANT" == "unrecorded" ]]; then
  echo "run-suite: WARNING: QUANT is unrecorded." >&2
  echo "           The endpoint cannot tell you which file is behind the alias, and a benchmark" >&2
  echo "           result without the quantisation is not reproducible. Set QUANT=Q4_K_M or" >&2
  echo "           whatever you actually served." >&2
fi

mkdir -p "$OUT_DIR"

ENDPOINT="$BASE_URL"
if [[ "$BACKEND" == "local-chat-completions" ]]; then
  ENDPOINT="${BASE_URL%/}/chat/completions"
else
  ENDPOINT="${BASE_URL%/}/completions"
fi

MODEL_ARGS="model=$ALIAS,base_url=$ENDPOINT,num_concurrent=$CONCURRENCY,max_retries=3,tokenized_requests=False"
if [[ -n "${TOKENIZER:-}" ]]; then
  MODEL_ARGS="$MODEL_ARGS,tokenizer=$TOKENIZER"
fi
if [[ -n "${API_KEY:-}" ]]; then
  export OPENAI_API_KEY="$API_KEY"
fi

echo "==> Suite: $TASKS"
echo "    alias:         $ALIAS"
echo "    backend:       $BACKEND -> $ENDPOINT"
echo "    quantisation:  $QUANT"
echo "    shots:         ${NUM_FEWSHOT:-the default in the task file}"
echo "    chat template: $CHAT_TEMPLATE"
echo "    gen_kwargs:    ${GEN_KWARGS:-the settings in the task file}"
echo "    runs:          $RUNS"
[[ -n "${LIMIT:-}" ]] && echo "    limit:         $LIMIT  (a limited run is not comparable with a published full run)"
echo

for run in $(seq 1 "$RUNS"); do
  seed=$(( run * 1000 + 7 ))
  echo "==> Run $run of $RUNS, seed $seed"

  ARGS=(
    --model "$BACKEND"
    --model_args "$MODEL_ARGS"
    --tasks "$TASKS"
    --batch_size 1
    --seed "$seed"
    --log_samples
    --output_path "$OUT_DIR/run-$run"
  )
  # Left unset, the shot count in the task file applies. Overriding it is a deliberate act
  # that has to appear in the report, which is why it is a variable rather than a default.
  if [[ -n "$NUM_FEWSHOT" ]]; then
    ARGS+=(--num_fewshot "$NUM_FEWSHOT")
  fi
  if [[ "$CHAT_TEMPLATE" == "yes" ]]; then
    ARGS+=(--apply_chat_template)
    if [[ -n "$NUM_FEWSHOT" && "$NUM_FEWSHOT" -gt 0 ]]; then
      ARGS+=(--fewshot_as_multiturn)
    fi
  fi
  if [[ -n "${LIMIT:-}" ]]; then
    ARGS+=(--limit "$LIMIT")
  fi
  # Generation settings override what the task file asked for. Setting them is how you turn a
  # greedy task into a sampled one, which is the third fault in the challenge.
  if [[ -n "${GEN_KWARGS:-}" ]]; then
    ARGS+=(--gen_kwargs "$GEN_KWARGS")
  fi

  lm_eval "${ARGS[@]}"
done

# The harness records everything it controls. These are the things it cannot see: which file the
# endpoint loaded, which engine is behind it, and which version. Without them the numbers above
# describe an alias rather than a model.
cat >"$OUT_DIR/run-context.json" <<JSON
{
  "alias": "$ALIAS",
  "tasks": "$TASKS",
  "backend": "$BACKEND",
  "endpoint": "$ENDPOINT",
  "quant": "$QUANT",
  "engine": "$ENGINE",
  "engine_version": "$ENGINE_VERSION",
  "num_fewshot": "${NUM_FEWSHOT:-task default}",
  "chat_template": "$CHAT_TEMPLATE",
  "gen_kwargs": "${GEN_KWARGS:-task default}",
  "runs": $RUNS,
  "limit": "${LIMIT:-none}",
  "date": "$(date +%Y-%m-%d)"
}
JSON

echo
echo "==> Done. $RUNS run(s) under $OUT_DIR/, settings in $OUT_DIR/run-context.json"
echo "    Summarise with:"
echo "      python3 summarise-results.py --results $OUT_DIR --labbook labbook.md"