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

GPU="0"
RELATION="bathroom_toilet"
BASE_MODEL="llava-hf/llava-1.5-7b-hf"
LORA_DIR=""
OUT_DIR=""
MAX_SAMPLES="0"
MAX_NEW_TOKENS="300"
DTYPE="float16"
PROMPT=""
PROMPTS=()
MENTION_ONLY="1"
LORA_ONLY="0"
BATCH_SIZE="1"
PROBE_BATCH_SIZE="0"
PPL_BATCH_SIZE="0"
ATTN_IMPL="eager"
MASTER_PORT="29500"
LOG=""

usage() {
  cat <<'EOF'
Usage:
  experiment/scripts/validate/run_eval_adv_lora_val.sh -l <lora_dir> [-g <gpu_ids>] [--full]

Required:
  -l <lora_dir>        LoRA adapter directory (e.g. adv_outputs/.../step_500)

Common:
  -g <gpu_ids>         GPU id(s) (default: 0). Examples: 0 or 0,1,2,3. Uses CUDA_VISIBLE_DEVICES=<gpu_ids>.
  -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 directory for captions.json / metrics.json (default: lora_dir)
  --max-samples <n>    Cap val examples (default: 0 = all)
  --log <path>         Tee stdout/stderr to file

Throughput:
  --batch-size <n>       Images per generate/forward call for captions (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>     Attention implementation: eager|sdpa|flash_attention_2 (default: eager)

Mention-only vs full:
  (default) mention-only: caption keyword mention rates only
  --full               Run probes + perplexity too (requires --sae-checkpoint and --probes-path)
  --lora-only          Skip base-model inference; run LoRA forward passes only

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

Other:
  --prompt <text>          Single prompt (overrides defaults)
  --prompts <p1> <p2> ...  Multiple prompts (space-separated; omit for both defaults)
  --dtype <float16|bfloat16>
  --max-new-tokens <n>
EOF
}

SAE_CHECKPOINT=""
PROBES_PATH=""

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

ARGS=()
while [[ $# -gt 0 ]]; do
  case "$1" in
    -g) GPU="$2"; shift 2 ;;
    -r) RELATION="$2"; shift 2 ;;
    -b) BASE_MODEL="$2"; shift 2 ;;
    -l) LORA_DIR="$2"; shift 2 ;;
    -o) OUT_DIR="$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 ;;
    --prompts) shift; while [[ $# -gt 0 && "$1" != -* ]]; do PROMPTS+=("$1"); shift; done ;;
    --log) LOG="$2"; shift 2 ;;
    --full) MENTION_ONLY="0"; shift 1 ;;
    --lora-only) LORA_ONLY="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 ;;
    --master-port) MASTER_PORT="$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 "${LORA_DIR}" ]]; then
  echo "Missing -l <lora_dir>" >&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 "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}"

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

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

CMD=("${LAUNCH[@]}"
  --relation "${RELATION}"
  --base_model "${BASE_MODEL}"
  --lora_dir "${LORA_DIR}"
  --max_samples "${MAX_SAMPLES}"
  --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 [[ ${#PROMPTS[@]} -gt 0 ]]; then
  CMD+=(--prompts "${PROMPTS[@]}")
elif [[ -n "${PROMPT}" ]]; then
  CMD+=(--prompt "${PROMPT}")
fi
# else: Python uses DEFAULT_EVAL_PROMPTS

if [[ "${LORA_ONLY}" == "1" ]]; then
  CMD+=(--lora_only)
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