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

export LD_LIBRARY_PATH=/data/data/com.termux/files/usr/lib/python3.13/site-packages/lib:/data/data/com.termux/files/usr/lib:$LD_LIBRARY_PATH

# Usage: ./bitnet_benchmark.sh /pfad/zum/model.gguf
VK_ICD_FILENAMES="$PREFIX/share/vulkan/icd.d/freedreno_icd.aarch64.json"
KMP_AFFINITY="disabled"
EXECUTABLE_DISABLE_MTE="1"
MODEL="${1:-models/quantized_q4km.gguf}"
PROMPT="${2:-Hallo BitNet}"
MAX_TOKENS="${3:-64}"
WARMUPS="${4:-3}"
RUNS="${5:-5}"
THREADS="${6:-$(nproc)}"
N_CTX="${7:-512}"

LOGDIR="${HOME}/bitnet_benchmark_logs"
mkdir -p "$LOGDIR"

echo "Model: $MODEL"
echo "Prompt: $PROMPT"
echo "Max new tokens: $MAX_TOKENS"
echo "Warmups: $WARMUPS, Runs: $RUNS, Threads: $THREADS, n_ctx: $N_CTX"
echo "Logs: $LOGDIR"
echo

# Helper to run inference and measure seconds
run_once() {
  local outlog="$1"
  local start end elapsed
  start=$(date +%s.%N)
  python3 ~/BitNet/run_inference.py -m "$MODEL" -t "$THREADS" \
    -c "$N_CTX" -n "$MAX_TOKENS" -p "$PROMPT" --no-stream \
    >"$outlog" 2>&1 || true
  end=$(date +%s.%N)
  elapsed=$(echo "$end - $start" | bc -l)
  printf "%.6f\n" "$elapsed"
}

# 1) Measure cold load + first token (single run)
echo "Measuring cold load + first token..."
COLDLOG="$LOGDIR/cold_run.log"
cold_time=$(run_once "$COLDLOG")
echo "Cold run time: $cold_time s"
echo

# 2) Warmups
echo "Warmup runs ($WARMUPS)..."
for i in $(seq 1 $WARMUPS); do
  wlog="$LOGDIR/warmup_${i}.log"
  t=$(run_once "$wlog")
  echo " Warmup $i: ${t}s"
done
echo

# 3) Timed runs
echo "Timed runs ($RUNS)..."
times_file="$LOGDIR/times.txt"
: > "$times_file"
for i in $(seq 1 $RUNS); do
  rlog="$LOGDIR/run_${i}.log"
  t=$(run_once "$rlog")
  echo "$t" | tee -a "$times_file"
  echo " Run $i: ${t}s"
done

# 4) Stats
python3 - <<PY
import statistics,sys
times = [float(x.strip()) for x in open("$times_file") if x.strip()]
if not times:
    print("No timed runs recorded")
    sys.exit(1)
median = statistics.median(times)
mean = statistics.mean(times)
minv = min(times)
maxv = max(times)
tokens = $MAX_TOKENS
print("---- Summary ----")
print(f"Runs: {len(times)}")
print(f"Median time per run: {median:.4f} s")
print(f"Mean time per run: {mean:.4f} s")
print(f"Min: {minv:.4f} s  Max: {maxv:.4f} s")
print(f"Tokens generated per run: {tokens}")
print(f"Tokens per second (median): {tokens/median:.2f}")
print(f"Tokens per second (mean): {tokens/mean:.2f}")
PY

echo
echo "Detaillogs liegen in $LOGDIR"
echo "Fertig"