#!/usr/bin/env bash
# Macstripe Lab — M4 Mac Mini local LLM benchmark (Ollama)
# Reproduces tok/s numbers cited in:
#   m4-mac-mini-7b-vs-14b-zhenshi-tiyan-2026.html
#
# Requirements: macOS + Apple Silicon, Ollama >= 0.6, curl, python3
# Usage:
#   chmod +x benchmark-m4-mac-mini-ollama.sh
#   ./benchmark-m4-mac-mini-ollama.sh llama3.1:8b
#   RUNS=5 PROMPT_TOKENS=512 GEN_TOKENS=256 ./benchmark-m4-mac-mini-ollama.sh qwen2.5:7b

set -euo pipefail

MODEL="${1:-llama3.1:8b}"
RUNS="${RUNS:-5}"
PROMPT_TOKENS="${PROMPT_TOKENS:-512}"
GEN_TOKENS="${GEN_TOKENS:-256}"
OLLAMA_HOST="${OLLAMA_HOST:-http://127.0.0.1:11434}"

echo "=== Macstripe M4 Mac Mini Ollama benchmark ==="
echo "model:          $MODEL"
echo "runs:           $RUNS"
echo "prompt tokens:  ~$PROMPT_TOKENS (fixed filler text)"
echo "gen tokens:     $GEN_TOKENS"
echo "ollama:         $OLLAMA_HOST"
echo ""

# Record idle memory / swap baseline
echo "--- idle baseline (before load) ---"
vm_stat | head -8
echo ""

# Fixed prompt (~512 tokens when repeated)
FILLER="Summarize unified memory architecture on Apple Silicon for local LLM inference. "
PROMPT=""
while [ "$(printf '%s' "$PROMPT" | wc -c | tr -d ' ')" -lt 1200 ]; do
  PROMPT="${PROMPT}${FILLER}"
done

pull_json=$(curl -sf "${OLLAMA_HOST}/api/pull" -d "{\"name\":\"${MODEL}\"}" || true)
if [ -z "$pull_json" ]; then
  echo "Warning: could not verify model pull; ensure ollama serve is running." >&2
fi

tok_rates=()
for i in $(seq 1 "$RUNS"); do
  echo "--- run $i / $RUNS ---"
  start_ms=$(python3 -c 'import time; print(int(time.time()*1000))')
  resp=$(curl -sf "${OLLAMA_HOST}/api/generate" -d "$(python3 - <<PY
import json
print(json.dumps({
  "model": "$MODEL",
  "prompt": """$PROMPT""",
  "stream": False,
  "options": {"num_predict": $GEN_TOKENS, "temperature": 0.2}
}))
PY
)")
  end_ms=$(python3 -c 'import time; print(int(time.time()*1000))')
  eval_count=$(printf '%s' "$resp" | python3 -c 'import sys,json; print(json.load(sys.stdin).get("eval_count",0))')
  elapsed_s=$(python3 -c "print(max(0.001, ($end_ms - $start_ms)/1000))")
  rate=$(python3 -c "print(round($eval_count / $elapsed_s, 1))")
  tok_rates+=("$rate")
  echo "  eval_count=$eval_count  elapsed=${elapsed_s}s  tok/s=$rate"
  sleep 2
done

echo ""
echo "--- swap after runs ---"
vm_stat | head -8
echo ""

median=$(printf '%s\n' "${tok_rates[@]}" | sort -n | python3 -c "
import sys
vals=[float(x) for x in sys.stdin if x.strip()]
vals.sort()
n=len(vals)
print(vals[n//2] if n else 0)
")

stats=$(printf '%s\n' "${tok_rates[@]}" | python3 -c "
import sys
vals=[float(x) for x in sys.stdin if x.strip()]
vals.sort()
n=len(vals)
if not n:
    print('mean=0 p50=0 p90=0')
else:
    mean=sum(vals)/n
    p50=vals[n//2]
    p90=vals[int(n*0.9)] if n>1 else vals[-1]
    print(f'mean={mean:.1f} p50={p50:.1f} p90={p90:.1f}')
")

echo "=== results ==="
echo "tok/s per run: ${tok_rates[*]}"
echo "median tok/s:  $median"
echo "$stats"
echo ""
echo "Paste into your lab notes with: sw_vers, sysctl hw.memsize, ollama --version"