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

export CUDA_VISIBLE_DEVICES=0,1,2,6

seed=${1:-42}
model_path=${2:-"liuhaotian/llava-v1.5-13b"}

cd_alpha=${3:-1}
cd_beta=${4:-0.5}
noise_step=${5:-500}

sigma=0.6
alpha=0.6
beta=0.8
start_layer=10
end_layer=$((start_layer + 30))

Q_FILE="/path/to/AMBER/query/query_all.json"
IMG_DIR="/path/to/AMBER/image"
DEPTH_DIR="/path/to/AMBER/depth"

OUT_DIR="./output/AMBER"
mkdir -p "${OUT_DIR}"

GEN_PY="./eval/object_hallucination_vqa_llava_amber2.py"
EVAL_PY="./eval/eval_amber.py"

out_name () {
  local tag="$1"
  echo "llava15_amber_answers_seed${seed}_${tag}_alpha${alpha}_beta${beta}_sigma${sigma}_start${start_layer}_end${end_layer}.json"
}

run_one () {
  local tag="$1"
  shift
  local out
  out=$(out_name "${tag}")

  echo ""
  echo "===== ${tag} ====="
  python -u "${GEN_PY}" \
    --model-path "${model_path}" \
    --question-file "${Q_FILE}" \
    --image-folder "${IMG_DIR}" \
    --answers-file "${OUT_DIR}/${out}" \
    --seed "${seed}" \
    --dscr-alpha "${alpha}" \
    --dscr-beta "${beta}" \
    --dscr-sigma "${sigma}" \
    --dscr-start-layer "${start_layer}" \
    --dscr-end-layer "${end_layer}" \
    --cd_alpha "${cd_alpha}" \
    --cd_beta "${cd_beta}" \
    --noise_step "${noise_step}" \
    "$@"

  python -u "${EVAL_PY}" \
    --inference_data "${OUT_DIR}/${out}" \
    --file-name "${out}" \
    --evaluation_type a
}

run_one "baseline"

run_one "vcd" \
   --use_vcd

run_one "dscr" \
  --use-dscr \
  --depth-folder "${DEPTH_DIR}" \
  --dscr-lambda 0.1 \
  --dscr-keep 1.0
