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

export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-1}"

repo_root="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
workdir="$repo_root/compare_model/ReplaceMe"
default_config="$workdir/examples/llama2_7b_replaceme_lstsq_skip16_common.yaml"
config_path="${CONFIG_PATH:-$default_config}"
target_tokens="${COMMON_TARGET_TOKENS_OVERRIDE:-}"
passthrough_args=()

while [[ $# -gt 0 ]]; do
  case "$1" in
    --target_tokens)
      if [[ $# -lt 2 ]]; then
        echo "error: --target_tokens requires a value" >&2
        exit 1
      fi
      target_tokens="$2"
      shift 2
      ;;
    --target_tokens=*)
      target_tokens="${1#*=}"
      shift
      ;;
    *)
      passthrough_args+=("$1")
      shift
      ;;
  esac
done

mkdir -p "$repo_root/results/llama_7b_replaceme_common_16"
git_commit="unknown"
if git -C "$repo_root" rev-parse --is-inside-work-tree >/dev/null 2>&1; then
  git_commit=$(git -C "$repo_root" rev-parse HEAD)
fi
{
  echo "git_commit=$git_commit"
  echo "config_path=$config_path"
  echo "COMMON_TARGET_TOKENS_OVERRIDE=${target_tokens:-4500000}"
  echo "command:"
  printf '%q ' python "$repo_root/compare_model/ReplaceMe/run_replaceme.py" --config "$config_path"
  if [[ -n "$target_tokens" ]]; then
    printf '%q ' "# COMMON_TARGET_TOKENS_OVERRIDE=$target_tokens"
  fi
  if [[ ${#passthrough_args[@]} -gt 0 ]]; then
    printf '%q ' "${passthrough_args[@]}"
  fi
  echo
} > "$repo_root/results/llama_7b_replaceme_common_16/run_args.txt"

cd "$workdir"
COMMON_TARGET_TOKENS_OVERRIDE="${target_tokens:-${COMMON_TARGET_TOKENS_OVERRIDE:-}}" \
PYTHONPATH="$workdir:$repo_root${PYTHONPATH:+:$PYTHONPATH}" \
python run_replaceme.py --config "$config_path" "${passthrough_args[@]}"