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

SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
PACKAGE_ROOT="$(cd -- "$SCRIPT_DIR/.." && pwd)"
WORKSPACE_ROOT="$(cd -- "$PACKAGE_ROOT/.." && pwd)"
BENCH="${BENCH:-$WORKSPACE_ROOT/tools-llama.cpp-k2-horizon/build-sm86/bin/llama-bench}"
MODEL_DIR="${MODEL_DIR:-$PACKAGE_ROOT}"
THREADS="${THREADS:-64}"
REPETITIONS="${REPETITIONS:-3}"
OUTDIR="${OUTDIR:-$WORKSPACE_ROOT/reports/k2-horizon-0.9b/benchmark-gpu1}"
RAW="$OUTDIR/llama-bench-raw"

mkdir -p "$RAW"
[[ -x "$BENCH" ]] || { echo "Missing BENCH: $BENCH" >&2; exit 1; }

export CUDA_VISIBLE_DEVICES=1
{
  echo "gpu_physical=1"
  echo "benchmark=llama-bench"
  echo "prompt_tokens=256"
  echo "generation_tokens=64"
  echo "batch_size=512"
  echo "ubatch_size=512"
  echo "threads=$THREADS"
  echo "repetitions=$REPETITIONS"
  nvidia-smi --query-gpu=index,name,memory.total,memory.used,utilization.gpu --format=csv
} > "$OUTDIR/environment.txt"

models=(
  "Q8_0|$MODEL_DIR/K2-Horizon-0.9B-Q8_0.gguf"
  "Q6_K|$MODEL_DIR/K2-Horizon-0.9B-Q6_K.gguf"
  "Q5_K_M|$MODEL_DIR/K2-Horizon-0.9B-Q5_K_M.gguf"
  "Q4_K_M|$MODEL_DIR/K2-Horizon-0.9B-Q4_K_M.gguf"
  "Q3_K_M|$MODEL_DIR/K2-Horizon-0.9B-Q3_K_M.gguf"
  "Q2_K|$MODEL_DIR/K2-Horizon-0.9B-Q2_K.gguf"
  "IQ2_XS|$MODEL_DIR/K2-Horizon-0.9B-IQ2_XS.gguf"
  "IQ1_M|$MODEL_DIR/K2-Horizon-0.9B-IQ1_M.gguf"
  "Q1_0|$MODEL_DIR/K2-Horizon-0.9B-Q1_0.gguf"
)

for item in "${models[@]}"; do
  label="${item%%|*}"
  model="${item#*|}"
  [[ -s "$model" ]] || { echo "Missing model: $model" >&2; exit 1; }
  echo "benchmarking GPU1 $label"
  timeout 900 "$BENCH" \
    -m "$model" \
    -p 256 \
    -n 64 \
    -r "$REPETITIONS" \
    -b 512 \
    -ub 512 \
    -t "$THREADS" \
    -ngl 99 \
    --device CUDA0 \
    -o csv \
    > "$RAW/$label.csv" \
    2> "$RAW/$label.stderr"
done

LLAMA_BENCH_REPETITIONS="$REPETITIONS" python3 \
  "$WORKSPACE_ROOT/reproducibility/summarize_bench.py" \
  "$RAW" "$OUTDIR/speed-summary.tsv"

echo "benchmark artifacts: $OUTDIR"