#!/bin/sh
# Compare Gem's real-weight logits with the adjacent llama.cpp reference runtime.
set -eu

usage() {
    echo "usage: $0 MODEL.gguf [OUTPUT.json] [PROMPT]" >&2
    exit 2
}

[ "$#" -ge 1 ] && [ "$#" -le 3 ] || usage

model=$1
output=${2:-}
prompt=${3:-Hello}
root=$(CDPATH='' cd -- "$(dirname -- "$0")/.." && pwd)
llama_debug=${GEM_LLAMA_DEBUG:-$root/../qwen38-llama-bench/llama.cpp/build-gem-metal/bin/llama-debug}
gem_logits=${GEM_LOGITS:-$root/build/gem-logits}
gem_tokenize=${GEM_TOKENIZE:-$root/build/gem-tokenize}
gem_compare=${GEM_LOGIT_COMPARE:-$root/build/gem-logit-compare}

[ -r "$model" ] || { echo "gem: model is not readable: $model" >&2; exit 1; }
[ -x "$llama_debug" ] || { echo "gem: llama-debug oracle is not executable: $llama_debug" >&2; exit 1; }
[ -x "$gem_logits" ] || { echo "gem: gem-logits is not executable: $gem_logits" >&2; exit 1; }
[ -x "$gem_tokenize" ] || { echo "gem: gem-tokenize is not executable: $gem_tokenize" >&2; exit 1; }
[ -x "$gem_compare" ] || { echo "gem: gem-logit-compare is not executable: $gem_compare" >&2; exit 1; }
command -v jq >/dev/null 2>&1 || { echo "gem: jq is required" >&2; exit 1; }

scratch=$(mktemp -d /private/tmp/gem-real-logit-oracle.XXXXXX)
cleanup() {
    case "$scratch" in
        /private/tmp/gem-real-logit-oracle.*) rm -rf -- "$scratch" ;;
    esac
}
trap cleanup EXIT HUP INT TERM

"$llama_debug" -m "$model" -p "$prompt" -c 64 -ngl all \
    -ctk f32 -ctv f32 --flash-attn off --save-logits \
    --logits-output-dir "$scratch" --log-disable

oracle_logits=$(find "$scratch" -maxdepth 1 -type f -name 'llamacpp-*.bin' \
    ! -name '*-tokens.bin' -print | head -n 1)
oracle_prompt=$(find "$scratch" -maxdepth 1 -type f -name 'llamacpp-*-prompt.txt' \
    -print | head -n 1)
[ -n "$oracle_logits" ] && [ -n "$oracle_prompt" ] || {
    echo "gem: llama-debug did not produce the expected oracle files" >&2
    exit 1
}

oracle_tokens=$(awk -F ': ' '/^token ids:/ {gsub(/, /, ",", $2); print "[" $2 "]"}' \
    "$oracle_prompt" | jq -c .)
native_tokens=$("$gem_tokenize" "$model" "$prompt" | jq -c .)
[ "$oracle_tokens" = "$native_tokens" ] || {
    echo "gem: tokenizer mismatch: oracle=$oracle_tokens native=$native_tokens" >&2
    exit 1
}

native_logits=$scratch/gem-native.bin
"$gem_logits" --model "$model" --prompt "$prompt" --context 64 --output "$native_logits"
comparison=$("$gem_compare" "$oracle_logits" "$native_logits")

record=$(jq -n --arg model "$model" --arg prompt "$prompt" \
    --argjson tokens "$native_tokens" --argjson comparison "$comparison" '
    {
      schema_version: 1,
      workload: "real-weight-first-token-logit-oracle",
      model: $model,
      prompt: $prompt,
      context: 64,
      kv_type: "f32",
      reference: "adjacent-llama.cpp-debug-metal-flash-off",
      token_ids: $tokens,
      comparison: $comparison,
      pass: ($comparison.finite_baseline == $comparison.count and
             $comparison.finite_candidate == $comparison.count and
             $comparison.same_argmax and $comparison.same_top10_set and
             $comparison.correlation >= 0.999 and $comparison.nmse <= 0.001)
    }')

printf '%s\n' "$record"
if [ -n "$output" ]; then
    printf '%s\n' "$record" > "$output"
fi

[ "$(printf '%s\n' "$record" | jq -r '.pass')" = true ]