#!/usr/bin/env bash
# Purpose: measure every quantisation in a directory the same way - speed and memory with
#          llama-bench, distributional damage with kl-divergence.py against the full-precision
#          reference, and task behaviour with the Part 10 harness - writing one result file per
#          quantisation so that measure-quants.py can assemble the comparison table
# Platform: all (spark, strix, mac, nvidia); GGUF rows only, one server at a time
# Minimum memory: 12 GB. The reference and the quantisations are served one after the other, not
#          together, so the floor is set by the largest single file plus its KV cache
# Assumes: llama-bench and llama-server on PATH or under LLAMA_BIN, curl, python3,
#          kl-divergence.py next to this script, and Part 10's run-eval.py and task file
#          reachable through EVAL_DIR and TASKS
#
# Usage: bash measure-quants.sh QUANT_DIR REFERENCE_GGUF
#   QUANT_DIR       directory of quantised GGUF files from quantise-five-ways.sh
#   REFERENCE_GGUF  the full-precision file every measurement is compared against
#
# Environment: LLAMA_BIN, PORT (default 8099), CTX (default 4096), NGL (default 99),
#              TEXT (calibration text for the divergence walk), POSITIONS (default 64),
#              EVAL_DIR (Part 10 lab directory), TASKS (Part 10 task file),
#              RESULTS (default results/), LABBOOK (default labbook.md), SKIP_EVAL=1
set -euo pipefail

QUANT_DIR="${1:-}"
REFERENCE="${2:-}"
PORT="${PORT:-8099}"
CTX="${CTX:-4096}"
NGL="${NGL:-99}"
POSITIONS="${POSITIONS:-64}"
RESULTS="${RESULTS:-results}"
LABBOOK="${LABBOOK:-labbook.md}"
HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
SERVER_PID=""

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

cleanup() {
  if [[ -n "$SERVER_PID" ]] && kill -0 "$SERVER_PID" 2>/dev/null; then
    kill "$SERVER_PID" 2>/dev/null || true
    wait "$SERVER_PID" 2>/dev/null || true
  fi
}
trap cleanup EXIT

[[ -n "$QUANT_DIR" && -n "$REFERENCE" ]] || die "usage: bash measure-quants.sh QUANT_DIR REFERENCE_GGUF"
[[ -d "$QUANT_DIR" ]] || die "$QUANT_DIR is not a directory"
[[ -f "$REFERENCE" ]] || die "$REFERENCE does not exist"
[[ -f "$HERE/kl-divergence.py" ]] || die "kl-divergence.py is not next to this script"
command -v curl >/dev/null || die "curl is not on PATH"
command -v python3 >/dev/null || die "python3 is not on PATH"

BENCH_BIN="${LLAMA_BIN:-}/llama-bench"
SERVER_BIN="${LLAMA_BIN:-}/llama-server"
if [[ -z "${LLAMA_BIN:-}" ]]; then
  BENCH_BIN="$(command -v llama-bench || true)"
  SERVER_BIN="$(command -v llama-server || true)"
fi
[[ -x "$BENCH_BIN" ]] || die "llama-bench not found; set LLAMA_BIN (built in Part 6)"
[[ -x "$SERVER_BIN" ]] || die "llama-server not found; set LLAMA_BIN (built in Part 6)"

TEXT="${TEXT:-}"
[[ -n "$TEXT" && -f "$TEXT" ]] || die "set TEXT to a plain-text file for the divergence walk (not your evaluation set)"

mkdir -p "$RESULTS"
BASE_URL="http://127.0.0.1:$PORT/v1"

start_server() {
  local model="$1" alias="$2"
  echo "    starting llama-server on port $PORT"
  "$SERVER_BIN" \
    --model "$model" \
    --alias "$alias" \
    --host 127.0.0.1 \
    --port "$PORT" \
    --ctx-size "$CTX" \
    --n-gpu-layers "$NGL" \
    --seed 0 \
    >"$RESULTS/server-$alias.log" 2>&1 &
  SERVER_PID=$!
  for _ in $(seq 1 180); do
    if curl -sf "http://127.0.0.1:$PORT/health" >/dev/null 2>&1; then
      return 0
    fi
    if ! kill -0 "$SERVER_PID" 2>/dev/null; then
      die "llama-server exited while loading $model; see $RESULTS/server-$alias.log"
    fi
    sleep 1
  done
  die "llama-server did not become healthy within 180s; see $RESULTS/server-$alias.log"
}

stop_server() {
  cleanup
  SERVER_PID=""
  sleep 2
}

run_task_set() {
  local alias="$1" label="$2"
  if [[ "${SKIP_EVAL:-0}" == "1" ]]; then
    echo "    task set skipped (SKIP_EVAL=1)"
    return 0
  fi
  if [[ -z "${EVAL_DIR:-}" || ! -f "${EVAL_DIR:-}/run-eval.py" ]]; then
    echo "    task set skipped: set EVAL_DIR to your Part 10 lab directory to include it"
    return 0
  fi
  local tasks="${TASKS:-$EVAL_DIR/tasks-template.json}"
  [[ -f "$tasks" ]] || die "task file $tasks does not exist"
  python3 "$EVAL_DIR/run-eval.py" \
    --base-url "$BASE_URL" \
    --model "$alias" \
    --quant "$label" \
    --engine llama.cpp \
    --tasks "$tasks" \
    --out "$RESULTS/eval-$label.json" \
    --notes "part-16 quantisation comparison"
}

# ---------------------------------------------------------------------------
# Pass one: the reference. Its distributions are saved to disk so that the
# quantisations can be measured against them without both models being resident.
# ---------------------------------------------------------------------------
echo "==> Reference: $(basename "$REFERENCE")"
start_server "$REFERENCE" reference
python3 "$HERE/kl-divergence.py" \
  --text "$TEXT" \
  --positions "$POSITIONS" \
  --base-url-ref "$BASE_URL" \
  --model-ref reference \
  --save-ref "$RESULTS/reference-logprobs.json"
run_task_set reference BF16
stop_server

# ---------------------------------------------------------------------------
# Pass two: every quantisation, one at a time, measured identically.
# ---------------------------------------------------------------------------
shopt -s nullglob
QUANTS=("$QUANT_DIR"/*.gguf)
shopt -u nullglob
[[ ${#QUANTS[@]} -gt 0 ]] || die "no .gguf files in $QUANT_DIR"

for quant in "${QUANTS[@]}"; do
  if [[ "$(cd "$(dirname "$quant")" && pwd)/$(basename "$quant")" == "$(cd "$(dirname "$REFERENCE")" && pwd)/$(basename "$REFERENCE")" ]]; then
    continue
  fi
  label="$(basename "$quant" .gguf)"
  echo
  echo "==> $label"

  echo "    llama-bench: prefill and decode, three repetitions"
  "$BENCH_BIN" \
    --model "$quant" \
    --n-prompt 512 \
    --n-gen 128 \
    --repetitions 3 \
    --n-gpu-layers "$NGL" \
    --output json \
    >"$RESULTS/bench-$label.json"

  start_server "$quant" quant
  echo "    kl-divergence against the saved reference distributions"
  python3 "$HERE/kl-divergence.py" \
    --text "$TEXT" \
    --positions "$POSITIONS" \
    --load-ref "$RESULTS/reference-logprobs.json" \
    --base-url-quant "$BASE_URL" \
    --model-quant quant \
    --quant-label "$label" \
    --out "$RESULTS/kld-$label.json" \
    --labbook "$LABBOOK"
  run_task_set quant "$label"
  stop_server
done

echo
echo "==> Every quantisation measured. Results in $RESULTS/"
echo "    Assemble the table with:"
echo "      python3 $HERE/measure-quants.py --results $RESULTS --labbook $LABBOOK"