#!/usr/bin/env bash
# Purpose: produce the GGUF half of the five-way comparison from one source checkpoint - a
#          full-precision reference plus Q8_0, Q5_K_M, Q4_K_M with an importance matrix and
#          Q4_K_M without one - so that every later measurement has the same starting point
# Platform: all (spark, strix, mac, nvidia); uses the llama.cpp build from Part 6
# Minimum memory: 12 GB; conversion holds one full-precision copy, quantisation streams
# Assumes: a llama.cpp checkout holding convert_hf_to_gguf.py (LLAMA_CPP), llama-quantize on
#          PATH or under LLAMA_BIN, an importance matrix from make-imatrix.sh for the fourth
#          file, and disk for roughly three times the full-precision size of the model
#
# Usage: bash quantise-five-ways.sh SOURCE OUTDIR [IMATRIX]
#   SOURCE   a Hugging Face model directory (safetensors) or an existing full-precision GGUF
#   OUTDIR   where the GGUF files are written
#   IMATRIX  the matrix from make-imatrix.sh; omit to skip the imatrix row
#
# Environment: LLAMA_CPP (default ~/llama.cpp), LLAMA_BIN, NAME (output prefix),
#              EXTRA_TYPE (a fifth GGUF type such as IQ4_XS, quantised with the matrix),
#              LABBOOK (default labbook.md)
set -euo pipefail

SOURCE="${1:-}"
OUTDIR="${2:-}"
IMATRIX="${3:-}"
LLAMA_CPP="${LLAMA_CPP:-$HOME/llama.cpp}"
LABBOOK="${LABBOOK:-labbook.md}"

die() { echo "quantise-five-ways: $*" >&2; exit 1; }

[[ -n "$SOURCE" && -n "$OUTDIR" ]] || die "usage: bash quantise-five-ways.sh SOURCE OUTDIR [IMATRIX]"
[[ -e "$SOURCE" ]] || die "$SOURCE does not exist"
command -v python3 >/dev/null || die "python3 is not on PATH"

QUANTIZE_BIN="${LLAMA_BIN:-}/llama-quantize"
if [[ -z "${LLAMA_BIN:-}" ]]; then
  QUANTIZE_BIN="$(command -v llama-quantize || true)"
fi
[[ -x "$QUANTIZE_BIN" ]] || die "llama-quantize not found; set LLAMA_BIN to the directory holding it (built in Part 6)"

mkdir -p "$OUTDIR"
NAME="${NAME:-$(basename "${SOURCE%.gguf}")}"
REFERENCE="$OUTDIR/$NAME-BF16.gguf"

# ---------------------------------------------------------------------------
# 1. The reference. Everything else in this part is measured against this file.
# ---------------------------------------------------------------------------
if [[ -f "$SOURCE" && "$SOURCE" == *.gguf ]]; then
  echo "==> 1/3 Using the supplied GGUF as the full-precision reference"
  REFERENCE="$SOURCE"
elif [[ -d "$SOURCE" ]]; then
  if [[ -f "$REFERENCE" ]]; then
    echo "==> 1/3 $REFERENCE already exists; reusing it"
  else
    [[ -f "$LLAMA_CPP/convert_hf_to_gguf.py" ]] || die "convert_hf_to_gguf.py not found under $LLAMA_CPP; set LLAMA_CPP"
    echo "==> 1/3 Converting $SOURCE to GGUF at bfloat16"
    echo "    Converting at full precision and quantising afterwards keeps information the"
    echo "    quantiser can use. Converting straight to a small type throws it away."
    python3 "$LLAMA_CPP/convert_hf_to_gguf.py" "$SOURCE" \
      --outfile "$REFERENCE" \
      --outtype bf16
  fi
else
  die "$SOURCE is neither a directory nor a .gguf file"
fi

# ---------------------------------------------------------------------------
# 2. The quantisations.
# ---------------------------------------------------------------------------
echo
echo "==> 2/3 Quantising"

quantise() {
  local type="$1" suffix="$2" use_imatrix="$3"
  local out="$OUTDIR/$NAME-$suffix.gguf"
  if [[ -f "$out" ]]; then
    echo "    $suffix already exists; skipping. Delete it to rebuild."
    return 0
  fi
  echo "    $suffix"
  if [[ "$use_imatrix" == "yes" ]]; then
    "$QUANTIZE_BIN" --imatrix "$IMATRIX" "$REFERENCE" "$out" "$type"
  else
    "$QUANTIZE_BIN" "$REFERENCE" "$out" "$type"
  fi
}

quantise q8_0 Q8_0 no
quantise q5_k_m Q5_K_M no
quantise q4_k_m Q4_K_M-noimat no

if [[ -n "$IMATRIX" ]]; then
  [[ -f "$IMATRIX" ]] || die "$IMATRIX does not exist; run make-imatrix.sh first or omit the argument"
  quantise q4_k_m Q4_K_M-imat yes
  if [[ -n "${EXTRA_TYPE:-}" ]]; then
    quantise "$(echo "$EXTRA_TYPE" | tr '[:upper:]' '[:lower:]')" "$EXTRA_TYPE-imat" yes
  fi
else
  echo "    no importance matrix supplied; the imatrix row will be missing from your table"
  if [[ -n "${EXTRA_TYPE:-}" ]]; then
    quantise "$(echo "$EXTRA_TYPE" | tr '[:upper:]' '[:lower:]')" "$EXTRA_TYPE" no
  fi
fi

# ---------------------------------------------------------------------------
# 3. Sizes, and the record.
# ---------------------------------------------------------------------------
echo
echo "==> 3/3 What you now have"
for f in "$OUTDIR/$NAME"-*.gguf "$REFERENCE"; do
  [[ -f "$f" ]] || continue
  printf '    %-44s %s bytes\n' "$(basename "$f")" "$(wc -c < "$f" | tr -d ' ')"
done | sort -u

python3 - "$LABBOOK" "$OUTDIR" "$NAME" "${IMATRIX:-none}" <<'PY'
import glob, json, os, platform, sys, time

labbook, outdir, name, imatrix = sys.argv[1:5]
files = sorted(set(glob.glob(os.path.join(outdir, f"{name}-*.gguf"))))
record = {
    "lab": "part-16/lab-quantise-five-ways-and-measure/quantise",
    "run_id": time.strftime("%Y%m%dT%H%M%S"),
    "name": name,
    "imatrix": os.path.basename(imatrix) if imatrix != "none" else None,
    "files": {os.path.basename(f): os.path.getsize(f) for f in files},
    "host": platform.platform(),
    "date": time.strftime("%Y-%m-%d"),
}
with open(labbook, "a", encoding="utf-8") as handle:
    handle.write(json.dumps(record) + "\n")
print(f"\nrecorded in {labbook}")
PY

echo
echo "    Keep the full-precision file until every measurement is finished. It is the reference"
echo "    every KL divergence in this lab is computed against, and it is not recoverable from"
echo "    the quantised ones."