#!/usr/bin/env bash

set -euo pipefail

# Validate that exactly 2 arguments are provided
if [ $# -ne 2 ]; then
  echo "Error: Exactly 2 arguments required."
  echo "Usage: $0 <llama_cpp_dir> <gguf>"
  echo "Example: $0 ~/code/llama.cpp BF16/Qwen3.5-35B-A3B-BF16-00001-of-00003.gguf"
  exit 1
fi

# Assign arguments to variables for clarity
LLAMA_CPP_DIR="$1"
GGUF_PATH="$2"

# Validate that the llama.cpp directory exists
if [ ! -d "$LLAMA_CPP_DIR" ]; then
  echo "Error: llama.cpp directory not found: $LLAMA_CPP_DIR"
  exit 1
fi


# Validate that the Python script exists
if [ ! -e "$GGUF_PATH" ]; then
  echo "Error: GGUF not found: $GGUF_PATH"
  exit 1
fi

CALIBRATION_DATASET_PATH=$HOME/calibration-data.txt

if [ ! -e $CALIBRATION_DATASET_PATH ]; then
  # download the calibration dataset
  curl -L https://gist.githubusercontent.com/ubergarm/edfeb3ff9c6ec8b49e88cdf627b0711a/raw/ba5b01b6960a86874592f5913e283746ff734483/ubergarm-imatrix-calibration-corpus-v02.txt > $CALIBRATION_DATASET_PATH
fi

# Get the directory where the script is located
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
IMATRIX_OUT_DIR="$(cd "$SCRIPT_DIR/.." && pwd)"
IMATRIX_PATH="$IMATRIX_OUT_DIR/imatrix.gguf"

if [ -e "$IMATRIX_PATH" ]; then
  echo "Error: imatrix already exists: $IMATRIX_PATH"
  exit 1
fi

$LLAMA_CPP_DIR/build/bin/llama-imatrix \
  -fit off \
  --model "$GGUF_PATH"\
  -f $CALIBRATION_DATASET_PATH \
  -o $IMATRIX_PATH \
  --ctx-size 512 \
  -ub 4096 -b 4096 \
  --no-mmap