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

GPU="0"
RELATION="bathroom_toilet"
BASE_MODEL="llava-hf/llava-1.5-7b-hf"
NULLU_CHECKPOINT=""
LOWEST_LAYER="16"
HIGHEST_LAYER="32"
OUT_DIR=""
MAX_SAMPLES="0"
MAX_NEW_TOKENS="300"
DTYPE="float16"
PROMPT="Describe this image."
MENTION_ONLY="1"
COMPARE_BASE="0"
BATCH_SIZE="1"
PROBE_BATCH_SIZE="0"
PPL_BATCH_SIZE="0"
ATTN_IMPL="eager"
SAE_CHECKPOINT=""
PROBES_PATH=""
LOG=""

usage() {
  cat <<'EOF'
Usage:
  scripts/validate/run_eval_nullu_val.sh -n <nullu_checkpoint> [-g <gpu_ids>] [--full]

Required:
  -n <nullu_checkpoint>   Path to Nullu edited-model directory (HF checkpoint).

Common:
  -g <gpu_ids>            GPU id(s) (default: 0). Examples: 0 or 0,1,2,3.
  -r <relation>           Relation key (default: bathroom_toilet)
  -b <base_model>         HF base model (default: llava-hf/llava-1.5-7b-hf)
  -o <output_dir>         Output dir for captions.json / metrics.json
                          (default: <nullu_checkpoint>/eval)
  --lowest-layer <n>      Inclusive lower bound of edited layer range (default: 16)
  --highest-layer <n>     Exclusive upper bound of edited layer range (default: 32)
  --max-samples <n>       Cap val examples per category (default: 0 = all)
  --log <path>            Tee stdout/stderr to file

Throughput:
  --batch-size <n>        Images per generate/forward call (default: 1; try 4–8 on A100)
  --probe-batch-size <n>  Batch size for probe eval (default: 0 = same as --batch-size)
  --ppl-batch-size <n>    Batch size for PPL eval   (default: 0 = same as --batch-size)
  --attn-impl <impl>      eager|sdpa|flash_attention_2 (default: eager)

Evaluation modes:
  (default) Mention-only: caption keyword mention rates only (Nullu-only, no base).
  --full                  Run probes + perplexity too (requires --sae-checkpoint and --probes-path).
  --compare-base          Also run base-model inference for comparison (two-pass).

Full-mode required:
  --sae-checkpoint <path>
  --probes-path <path>

Other:
  --prompt <text>
  --dtype <float16|bfloat16>
  --max-new-tokens <n>
EOF
}

if [[ $# -eq 0 ]]; then
  usage
  exit 2
fi

while [[ $# -gt 0 ]]; do
  case "$1" in
    -g) GPU="$2"; shift 2 ;;
    -r) RELATION="$2"; shift 2 ;;
    -b) BASE_MODEL="$2"; shift 2 ;;
    -n) NULLU_CHECKPOINT="$2"; shift 2 ;;
    -o) OUT_DIR="$2"; shift 2 ;;
    --lowest-layer) LOWEST_LAYER="$2"; shift 2 ;;
    --highest-layer) HIGHEST_LAYER="$2"; shift 2 ;;
    --max-samples) MAX_SAMPLES="$2"; shift 2 ;;
    --max-new-tokens) MAX_NEW_TOKENS="$2"; shift 2 ;;
    --dtype) DTYPE="$2"; shift 2 ;;
    --prompt) PROMPT="$2"; shift 2 ;;
    --log) LOG="$2"; shift 2 ;;
    --full) MENTION_ONLY="0"; shift 1 ;;
    --compare-base) COMPARE_BASE="1"; shift 1 ;;
    --batch-size) BATCH_SIZE="$2"; shift 2 ;;
    --probe-batch-size) PROBE_BATCH_SIZE="$2"; shift 2 ;;
    --ppl-batch-size) PPL_BATCH_SIZE="$2"; shift 2 ;;
    --attn-impl) ATTN_IMPL="$2"; shift 2 ;;
    --sae-checkpoint) SAE_CHECKPOINT="$2"; shift 2 ;;
    --probes-path) PROBES_PATH="$2"; shift 2 ;;
    -h|--help) usage; exit 0 ;;
    *) echo "Unknown arg: $1" >&2; usage; exit 2 ;;
  esac
done

if [[ -z "${NULLU_CHECKPOINT}" ]]; then
  echo "Missing -n <nullu_checkpoint>" >&2
  exit 2
fi

ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../../.." && pwd)"
cd "${ROOT}"

if [[ -d "venv" ]]; then
  # shellcheck disable=SC1091
  source venv/bin/activate
elif [[ -d "env" ]]; then
  # shellcheck disable=SC1091
  source env/bin/activate
fi

export PYTHONPATH="${ROOT}:${PYTHONPATH:-}"
export CUDA_VISIBLE_DEVICES="${GPU}"
export HF_HOME="${HF_HOME:-$HOME/.cache/huggingface}"

NPROC=1
if [[ "${GPU}" == *","* ]]; then
  NPROC=$(( $(tr -cd ',' <<<"${GPU}" | wc -c) + 1 ))
fi

if [[ "${NPROC}" -gt 1 ]]; then
  LAUNCH=(torchrun --nproc_per_node="${NPROC}" -m experiment.evaluation.eval_nullu_val)
else
  LAUNCH=(python -m experiment.evaluation.eval_nullu_val)
fi

CMD=("${LAUNCH[@]}"
  --relation "${RELATION}"
  --base_model "${BASE_MODEL}"
  --nullu_checkpoint "${NULLU_CHECKPOINT}"
  --lowest_layer "${LOWEST_LAYER}"
  --highest_layer "${HIGHEST_LAYER}"
  --max_samples "${MAX_SAMPLES}"
  --prompt "${PROMPT}"
  --max_new_tokens "${MAX_NEW_TOKENS}"
  --dtype "${DTYPE}"
  --batch_size "${BATCH_SIZE}"
  --probe_batch_size "${PROBE_BATCH_SIZE}"
  --ppl_batch_size "${PPL_BATCH_SIZE}"
  --attn_impl "${ATTN_IMPL}"
)

if [[ -n "${OUT_DIR}" ]]; then
  CMD+=(--output_dir "${OUT_DIR}")
fi

if [[ "${COMPARE_BASE}" == "1" ]]; then
  CMD+=(--compare_base)
fi

if [[ "${MENTION_ONLY}" == "1" ]]; then
  CMD+=(--mention_only)
else
  if [[ -z "${SAE_CHECKPOINT}" || -z "${PROBES_PATH}" ]]; then
    echo "--full requires --sae-checkpoint and --probes-path" >&2
    exit 2
  fi
  CMD+=(--sae_checkpoint "${SAE_CHECKPOINT}" --probes_path "${PROBES_PATH}")
fi

if [[ -n "${LOG}" ]]; then
  "${CMD[@]}" 2>&1 | tee -a "${LOG}"
else
  "${CMD[@]}"
fi