#!/usr/bin/env bash
set -euo pipefail
shopt -s nullglob

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

SEED="${SEED:-42}"
CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0,1,2,3,4,6}"
export CUDA_VISIBLE_DEVICES
export MASTER_PORT="${MASTER_PORT:-12345}"

MODELS="${MODELS:-llava15 qwen mplug}"
METHODS="${METHODS:-baseline dscr vcd opera vcd_dscr opera_dscr}"
POPE_SETS="${POPE_SETS:-coco aokvqa gqa}"

MODEL_PATH_LLAVA15="${MODEL_PATH_LLAVA15:-liuhaotian/llava-v1.5-13b}"
MODEL_PATH_QWEN="${MODEL_PATH_QWEN:-Qwen/Qwen-VL}"
MODEL_PATH_MPLUG="${MODEL_PATH_MPLUG:-MAGAer13/mplug-owl2-llama2-7b}"

POPE_ROOT="${POPE_ROOT:-/path/to/POPE}"
REPOPE_ROOT="${REPOPE_ROOT:-/path/to/REPOPE}"
OUT_ROOT="${OUT_ROOT:-${EXP_DIR}/output/POPE_all_models}"
mkdir -p "${OUT_ROOT}"

FORMAT="${FORMAT:-yn_format}"
MAX_NEW_TOKENS="${MAX_NEW_TOKENS:-2}"
TEMPERATURE="${TEMPERATURE:-0.0}"
TOP_P="${TOP_P:-1.0}"
TOP_K="${TOP_K:-0}"

# VCD
NOISE_STEP="${NOISE_STEP:-500}"
CD_ALPHA="${CD_ALPHA:-0.5}"
CD_BETA="${CD_BETA:-0.1}"

# OPERA
OPERA_BEAM="${OPERA_BEAM:-5}"
OPERA_SCALE_FACTOR="${OPERA_SCALE_FACTOR:-50.0}"
OPERA_THRESHOLD="${OPERA_THRESHOLD:-15}"
OPERA_NUM_ATTN_CANDIDATES="${OPERA_NUM_ATTN_CANDIDATES:-5}"
OPERA_PENALTY_WEIGHTS="${OPERA_PENALTY_WEIGHTS:-0.6}"

# DSCR defaults by model
LLAVA_DSCR_ALPHA="${LLAVA_DSCR_ALPHA:-0.6}"
LLAVA_DSCR_BETA="${LLAVA_DSCR_BETA:-0.8}"
LLAVA_DSCR_SIGMA="${LLAVA_DSCR_SIGMA:-0.6}"
LLAVA_DSCR_KEEP_RATIO="${LLAVA_DSCR_KEEP_RATIO:-1.0}"
LLAVA_DSCR_LAMBDA="${LLAVA_DSCR_LAMBDA:-1.0}"
LLAVA_DSCR_START_LAYER="${LLAVA_DSCR_START_LAYER:-10}"
LLAVA_DSCR_END_LAYER="${LLAVA_DSCR_END_LAYER:-40}"

QWEN_DSCR_ALPHA="${QWEN_DSCR_ALPHA:-0.6}"
QWEN_DSCR_BETA="${QWEN_DSCR_BETA:-0.8}"
QWEN_DSCR_SIGMA="${QWEN_DSCR_SIGMA:-0.6}"
QWEN_DSCR_KEEP_RATIO="${QWEN_DSCR_KEEP_RATIO:-1.0}"
QWEN_DSCR_LAMBDA="${QWEN_DSCR_LAMBDA:-1.0}"
QWEN_DSCR_START_LAYER="${QWEN_DSCR_START_LAYER:-16}"
QWEN_DSCR_END_LAYER="${QWEN_DSCR_END_LAYER:-24}"

MPLUG_DSCR_ALPHA="${MPLUG_DSCR_ALPHA:-0.6}"
MPLUG_DSCR_BETA="${MPLUG_DSCR_BETA:-0.55}"
MPLUG_DSCR_SIGMA="${MPLUG_DSCR_SIGMA:-0.24}"
MPLUG_DSCR_LAMBDA="${MPLUG_DSCR_LAMBDA:-1.0}"
MPLUG_DSCR_START_LAYER="${MPLUG_DSCR_START_LAYER:-12}"
MPLUG_DSCR_END_LAYER="${MPLUG_DSCR_END_LAYER:-20}"
DSCR_MODE="${DSCR_MODE:-key-only}"

function pope_paths() {
  local pope_set="$1"
  case "${pope_set}" in
    coco)
      GT_DIR="${POPE_ROOT}/coco"
      IMAGE_ROOT="${POPE_ROOT}/coco/val2014_modified"
      DEPTH_ROOT="${POPE_ROOT}/coco/POPE_coco_depth/npy"
      ;;
    aokvqa)
      GT_DIR="${POPE_ROOT}/aokvqa"
      IMAGE_ROOT="${POPE_ROOT}/aokvqa/AOK_val2014"
      DEPTH_ROOT="${POPE_ROOT}/aokvqa/AOK_coco_depth"
      ;;
    gqa)
      GT_DIR="${POPE_ROOT}/gqa"
      IMAGE_ROOT="${POPE_ROOT}/gqa/gqa_refine"
      DEPTH_ROOT="${POPE_ROOT}/gqa/gqa_depth"
      ;;
    *)
      echo "[ERR] Unknown POPE set: ${pope_set}"
      exit 1
      ;;
  esac
}

function dscr_mode_flags_mme() {
  DSCR_MODE_FLAG_MME=()
  DSCR_MODE_FLAG_MPLUG=()
  DSCR_MODE_TAG=""
  case "${DSCR_MODE}" in
    key-only)
      DSCR_MODE_FLAG_MME=(--dscr-key-only)
      DSCR_MODE_FLAG_MPLUG=(--key-only True --value-only False --key-value False)
      DSCR_MODE_TAG="key_only"
      ;;
    value-only)
      DSCR_MODE_FLAG_MME=(--dscr-value-only)
      DSCR_MODE_FLAG_MPLUG=(--key-only False --value-only True --key-value False)
      DSCR_MODE_TAG="value_only"
      ;;
    key-value)
      DSCR_MODE_FLAG_MME=(--dscr-key-value)
      DSCR_MODE_FLAG_MPLUG=(--key-only False --value-only False --key-value True)
      DSCR_MODE_TAG="key_value"
      ;;
    *)
      echo "[ERR] Unknown DSCR_MODE: ${DSCR_MODE}"
      exit 1
      ;;
  esac
}

function eval_one_file() {
  local gt_file="$1"
  local gen_file="$2"
  local base_name
  base_name="$(basename "${gen_file}")"
  python "${EXP_DIR}/eval/eval_pope.py" \
    --gt_files "${gt_file}" \
    --gen_files "${gen_file}" \
    --seed "${SEED}" \
    --file-name "${base_name}"
}

function eval_one_file_repope() {
  local pope_gt_file="$1"
  local gen_file="$2"
  local pope_set="$3"

  local repope_gt_file
  repope_gt_file="${pope_gt_file/${POPE_ROOT}/${REPOPE_ROOT}}"
  if [[ ! -f "${repope_gt_file}" ]]; then
    return 0
  fi

  local base_name
  base_name="repope_$(basename "${gen_file}")"
  echo "  [REPOPE] ${repope_gt_file}"
  python "${EXP_DIR}/eval/eval_pope.py" \
    --gt_files "${repope_gt_file}" \
    --gen_files "${gen_file}" \
    --seed "${SEED}" \
    --file-name "${base_name}"
}

function run_mme_model() {
  local eval_file="$1"
  local model_path="$2"
  local model_tag="$3"
  local pope_set="$4"
  local method="$5"
  local out_dir="$6"
  local dscr_alpha="$7"
  local dscr_beta="$8"
  local dscr_sigma="$9"
  local dscr_keep_ratio="${10}"
  local dscr_lambda="${11}"
  local dscr_start_layer="${12}"
  local dscr_end_layer="${13}"
  local method_tag="$method"
  local use_dscr_tag="nodscr"
  local extra_args=()

  case "${method}" in
    baseline)
      extra_args+=(
        --method baseline
        --temperature "${TEMPERATURE}"
        --top-p "${TOP_P}"
        --top-k "${TOP_K}"
      )
      ;;
    dscr)
      method_tag="baseline"
      use_dscr_tag="usedscr_${DSCR_MODE_TAG}"
      extra_args+=(
        --method baseline
        --use-dscr
        --depth-root "${DEPTH_ROOT}"
        --dscr-alpha "${dscr_alpha}"
        --dscr-beta "${dscr_beta}"
        --dscr-sigma "${dscr_sigma}"
        --dscr-keep-ratio "${dscr_keep_ratio}"
        --dscr-lambda "${dscr_lambda}"
        --dscr-start-layer "${dscr_start_layer}"
        --dscr-end-layer "${dscr_end_layer}"
        "${DSCR_MODE_FLAG_MME[@]}"
        --temperature 0.0
        --top-p 1.0
        --top-k 0
      )
      ;;
    vcd)
      method_tag="usevcd"
      extra_args+=(
        --method vcd
        --noise-step "${NOISE_STEP}"
        --cd-alpha "${CD_ALPHA}"
        --cd-beta "${CD_BETA}"
        --temperature 0.0
        --top-p "${TOP_P}"
        --top-k "${TOP_K}"
      )
      ;;
    opera)
      method_tag="useopera"
      extra_args+=(
        --method opera
        --beam "${OPERA_BEAM}"
        --scale-factor "${OPERA_SCALE_FACTOR}"
        --threshold "${OPERA_THRESHOLD}"
        --num-attn-candidates "${OPERA_NUM_ATTN_CANDIDATES}"
        --penalty-weights "${OPERA_PENALTY_WEIGHTS}"
        --temperature "${TEMPERATURE}"
        --top-p "${TOP_P}"
        --top-k "${TOP_K}"
      )
      ;;
    vcd_dscr)
      method_tag="usevcd_dscr"
      use_dscr_tag="usedscr_${DSCR_MODE_TAG}"
      extra_args+=(
        --method vcd
        --noise-step "${NOISE_STEP}"
        --cd-alpha "${CD_ALPHA}"
        --cd-beta "${CD_BETA}"
        --use-dscr
        --depth-root "${DEPTH_ROOT}"
        --dscr-alpha "${dscr_alpha}"
        --dscr-beta "${dscr_beta}"
        --dscr-sigma "${dscr_sigma}"
        --dscr-keep-ratio "${dscr_keep_ratio}"
        --dscr-lambda "${dscr_lambda}"
        --dscr-start-layer "${dscr_start_layer}"
        --dscr-end-layer "${dscr_end_layer}"
        "${DSCR_MODE_FLAG_MME[@]}"
        --temperature 0.0
        --top-p "${TOP_P}"
        --top-k "${TOP_K}"
      )
      ;;
    opera_dscr)
      method_tag="useopera_dscr"
      use_dscr_tag="usedscr_${DSCR_MODE_TAG}"
      extra_args+=(
        --method opera
        --beam "${OPERA_BEAM}"
        --scale-factor "${OPERA_SCALE_FACTOR}"
        --threshold "${OPERA_THRESHOLD}"
        --num-attn-candidates "${OPERA_NUM_ATTN_CANDIDATES}"
        --penalty-weights "${OPERA_PENALTY_WEIGHTS}"
        --use-dscr
        --depth-root "${DEPTH_ROOT}"
        --dscr-alpha "${dscr_alpha}"
        --dscr-beta "${dscr_beta}"
        --dscr-sigma "${dscr_sigma}"
        --dscr-keep-ratio "${dscr_keep_ratio}"
        --dscr-lambda "${dscr_lambda}"
        --dscr-start-layer "${dscr_start_layer}"
        --dscr-end-layer "${dscr_end_layer}"
        "${DSCR_MODE_FLAG_MME[@]}"
        --temperature "${TEMPERATURE}"
        --top-p "${TOP_P}"
        --top-k "${TOP_K}"
      )
      ;;
    *)
      echo "[ERR] Unknown method: ${method}"
      exit 1
      ;;
  esac

  run_name="${model_tag}_pope_${pope_set}_${method_tag}_${use_dscr_tag}_seed${SEED}_alpha${dscr_alpha}_beta${dscr_beta}_start${dscr_start_layer}_end${dscr_end_layer}"

  python "${eval_file}" \
    --model-path "${model_path}" \
    --gt-dir "${GT_DIR}" \
    --image-root "${IMAGE_ROOT}" \
    --out-root "${out_dir}" \
    --run-name "${run_name}" \
    --seed "${SEED}" \
    --format "${FORMAT}" \
    --max-new-tokens "${MAX_NEW_TOKENS}" \
    "${extra_args[@]}"

  local gen_file base_name set_name gt_json gt_jsonl gt_file
  for gen_file in "${out_dir}/${run_name}_"*.jsonl; do
    base_name="$(basename "${gen_file}")"
    set_name="${base_name#${run_name}_}"
    set_name="${set_name%.jsonl}"
    gt_json="${GT_DIR}/${set_name}.json"
    gt_jsonl="${GT_DIR}/${set_name}.jsonl"
    if [[ -f "${gt_json}" ]]; then
      gt_file="${gt_json}"
    elif [[ -f "${gt_jsonl}" ]]; then
      gt_file="${gt_jsonl}"
    else
      echo "[WARN] GT not found for ${base_name}"
      continue
    fi
    eval_one_file "${gt_file}" "${gen_file}"
    eval_one_file_repope "${gt_file}" "${gen_file}" "${pope_set}"
  done
}

function run_mplug_model() {
  local model_path="$1"
  local model_tag="$2"
  local pope_set="$3"
  local method="$4"
  local out_dir="$5"
  local dscr_alpha="$6"
  local dscr_beta="$7"
  local dscr_sigma="$8"
  local dscr_lambda="$9"
  local dscr_start_layer="${10}"
  local dscr_end_layer="${11}"
  local method_tag="${method}"
  local use_dscr_tag="nodscr"
  local extra_args=()

  case "${method}" in
    baseline)
      extra_args+=(
        --method baseline
        --temperature "${TEMPERATURE}"
        --top-p "${TOP_P}"
        --top-k "${TOP_K}"
      )
      ;;
    dscr)
      method_tag="baseline"
      use_dscr_tag="usedscr_${DSCR_MODE_TAG}"
      extra_args+=(
        --method baseline
        --use-dscr
        --depth-root "${DEPTH_ROOT}"
        --dscr-alpha "${dscr_alpha}"
        --dscr-beta "${dscr_beta}"
        --dscr-sigma "${dscr_sigma}"
        --dscr-lambda "${dscr_lambda}"
        --dscr-start-layer "${dscr_start_layer}"
        --dscr-end-layer "${dscr_end_layer}"
        "${DSCR_MODE_FLAG_MME[@]}"
        --temperature 0.0
        --top-p 1.0
        --top-k 0
      )
      ;;
    vcd)
      method_tag="usevcd"
      extra_args+=(
        --method vcd
        --noise-step "${NOISE_STEP}"
        --cd-alpha "${CD_ALPHA}"
        --cd-beta "${CD_BETA}"
        --temperature 0.0
        --top-p "${TOP_P}"
        --top-k "${TOP_K}"
      )
      ;;
    opera)
      method_tag="useopera"
      extra_args+=(
        --method opera
        --beam "${OPERA_BEAM}"
        --scale-factor "${OPERA_SCALE_FACTOR}"
        --threshold "${OPERA_THRESHOLD}"
        --num-attn-candidates "${OPERA_NUM_ATTN_CANDIDATES}"
        --penalty-weights "${OPERA_PENALTY_WEIGHTS}"
        --temperature "${TEMPERATURE}"
        --top-p "${TOP_P}"
        --top-k "${TOP_K}"
      )
      ;;
    vcd_dscr)
      method_tag="usevcd_dscr"
      use_dscr_tag="usedscr_${DSCR_MODE_TAG}"
      extra_args+=(
        --method vcd
        --noise-step "${NOISE_STEP}"
        --cd-alpha "${CD_ALPHA}"
        --cd-beta "${CD_BETA}"
        --use-dscr
        --depth-root "${DEPTH_ROOT}"
        --dscr-alpha "${dscr_alpha}"
        --dscr-beta "${dscr_beta}"
        --dscr-sigma "${dscr_sigma}"
        --dscr-lambda "${dscr_lambda}"
        --dscr-start-layer "${dscr_start_layer}"
        --dscr-end-layer "${dscr_end_layer}"
        "${DSCR_MODE_FLAG_MME[@]}"
        --temperature 0.0
        --top-p "${TOP_P}"
        --top-k "${TOP_K}"
      )
      ;;
    opera_dscr)
      method_tag="useopera_dscr"
      use_dscr_tag="usedscr_${DSCR_MODE_TAG}"
      extra_args+=(
        --method opera
        --beam "${OPERA_BEAM}"
        --scale-factor "${OPERA_SCALE_FACTOR}"
        --threshold "${OPERA_THRESHOLD}"
        --num-attn-candidates "${OPERA_NUM_ATTN_CANDIDATES}"
        --penalty-weights "${OPERA_PENALTY_WEIGHTS}"
        --use-dscr
        --depth-root "${DEPTH_ROOT}"
        --dscr-alpha "${dscr_alpha}"
        --dscr-beta "${dscr_beta}"
        --dscr-sigma "${dscr_sigma}"
        --dscr-lambda "${dscr_lambda}"
        --dscr-start-layer "${dscr_start_layer}"
        --dscr-end-layer "${dscr_end_layer}"
        "${DSCR_MODE_FLAG_MME[@]}"
        --temperature "${TEMPERATURE}"
        --top-p "${TOP_P}"
        --top-k "${TOP_K}"
      )
      ;;
    *)
      echo "[ERR] Unknown method: ${method}"
      exit 1
      ;;
  esac

  local run_name
  run_name="${model_tag}_pope_${pope_set}_${method_tag}_${use_dscr_tag}_seed${SEED}_alpha${dscr_alpha}_beta${dscr_beta}_start${dscr_start_layer}_end${dscr_end_layer}"

  python "${EXP_DIR}/eval/pope_mplug.py" \
    --model-path "${model_path}" \
    --gt-dir "${GT_DIR}" \
    --image-root "${IMAGE_ROOT}" \
    --out-root "${out_dir}" \
    --depth-root "${DEPTH_ROOT}" \
    --run-name "${run_name}" \
    --seed "${SEED}" \
    --format "${FORMAT}" \
    --max-new-tokens "${MAX_NEW_TOKENS}" \
    "${extra_args[@]}"

  local gen_file base_name set_name gt_json gt_jsonl gt_file
  for gen_file in "${out_dir}/${run_name}_"*.jsonl; do
    base_name="$(basename "${gen_file}")"
    set_name="${base_name#${run_name}_}"
    set_name="${set_name%.jsonl}"
    gt_json="${GT_DIR}/${set_name}.json"
    gt_jsonl="${GT_DIR}/${set_name}.jsonl"
    if [[ -f "${gt_json}" ]]; then
      gt_file="${gt_json}"
    elif [[ -f "${gt_jsonl}" ]]; then
      gt_file="${gt_jsonl}"
    else
      echo "[WARN] GT not found for ${base_name}"
      continue
    fi
    eval_one_file "${gt_file}" "${gen_file}"
    eval_one_file_repope "${gt_file}" "${gen_file}" "${pope_set}"
  done
}

dscr_mode_flags_mme

_resume_from="${RESUME_FROM:-}"
_resuming=true
if [[ -z "${_resume_from}" ]]; then
  _resuming=false
fi

for model_key in ${MODELS}; do
  case "${model_key}" in
    llava15)
      MODEL_DSCR_ALPHA="${LLAVA_DSCR_ALPHA}"
      MODEL_DSCR_BETA="${LLAVA_DSCR_BETA}"
      MODEL_DSCR_SIGMA="${LLAVA_DSCR_SIGMA}"
      MODEL_DSCR_KEEP_RATIO="${LLAVA_DSCR_KEEP_RATIO}"
      MODEL_DSCR_LAMBDA="${LLAVA_DSCR_LAMBDA}"
      MODEL_DSCR_START_LAYER="${LLAVA_DSCR_START_LAYER}"
      MODEL_DSCR_END_LAYER="${LLAVA_DSCR_END_LAYER}"
      ;;
    qwen)
      MODEL_DSCR_ALPHA="${QWEN_DSCR_ALPHA}"
      MODEL_DSCR_BETA="${QWEN_DSCR_BETA}"
      MODEL_DSCR_SIGMA="${QWEN_DSCR_SIGMA}"
      MODEL_DSCR_KEEP_RATIO="${QWEN_DSCR_KEEP_RATIO}"
      MODEL_DSCR_LAMBDA="${QWEN_DSCR_LAMBDA}"
      MODEL_DSCR_START_LAYER="${QWEN_DSCR_START_LAYER}"
      MODEL_DSCR_END_LAYER="${QWEN_DSCR_END_LAYER}"
      ;;
    mplug)
      MODEL_DSCR_ALPHA="${MPLUG_DSCR_ALPHA}"
      MODEL_DSCR_BETA="${MPLUG_DSCR_BETA}"
      MODEL_DSCR_SIGMA="${MPLUG_DSCR_SIGMA}"
      MODEL_DSCR_KEEP_RATIO="1.0"
      MODEL_DSCR_LAMBDA="${MPLUG_DSCR_LAMBDA}"
      MODEL_DSCR_START_LAYER="${MPLUG_DSCR_START_LAYER}"
      MODEL_DSCR_END_LAYER="${MPLUG_DSCR_END_LAYER}"
      ;;
    *)
      echo "[ERR] Unsupported model key: ${model_key}. Use llava15 qwen mplug."
      exit 1
      ;;
  esac

  for pope_set in ${POPE_SETS}; do
    pope_paths "${pope_set}"
    if [[ ! -d "${GT_DIR}" ]]; then
      echo "[WARN] Skip ${pope_set}: GT dir not found (${GT_DIR})"
      continue
    fi

    for method in ${METHODS}; do
      if [[ "${_resuming}" == true ]]; then
        if [[ "${_resume_from}" == "${model_key}:${pope_set}:${method}" ]]; then
          _resuming=false
          echo "[RESUME] Starting from model=${model_key} set=${pope_set} method=${method}"
        else
          echo "[SKIP] model=${model_key} set=${pope_set} method=${method}"
          continue
        fi
      fi

      out_dir="${OUT_ROOT}/${model_key}/${pope_set}/${method}"
      mkdir -p "${out_dir}"
      echo "===================================================================="
      echo "[RUN] model=${model_key} set=${pope_set} method=${method}"
      echo "===================================================================="

      case "${model_key}" in
        llava15)
          run_mme_model \
            "${EXP_DIR}/eval/mme_llava15.py" "${MODEL_PATH_LLAVA15}" "llava15" "${pope_set}" "${method}" "${out_dir}" \
            "${MODEL_DSCR_ALPHA}" "${MODEL_DSCR_BETA}" "${MODEL_DSCR_SIGMA}" "${MODEL_DSCR_KEEP_RATIO}" "${MODEL_DSCR_LAMBDA}" "${MODEL_DSCR_START_LAYER}" "${MODEL_DSCR_END_LAYER}"
          ;;
        qwen)
          run_mme_model \
            "${EXP_DIR}/eval/mme_qwen.py" "${MODEL_PATH_QWEN}" "qwen" "${pope_set}" "${method}" "${out_dir}" \
            "${MODEL_DSCR_ALPHA}" "${MODEL_DSCR_BETA}" "${MODEL_DSCR_SIGMA}" "${MODEL_DSCR_KEEP_RATIO}" "${MODEL_DSCR_LAMBDA}" "${MODEL_DSCR_START_LAYER}" "${MODEL_DSCR_END_LAYER}"
          ;;
        mplug)
          run_mplug_model \
            "${MODEL_PATH_MPLUG}" "mplug" "${pope_set}" "${method}" "${out_dir}" \
            "${MODEL_DSCR_ALPHA}" "${MODEL_DSCR_BETA}" "${MODEL_DSCR_SIGMA}" "${MODEL_DSCR_LAMBDA}" "${MODEL_DSCR_START_LAYER}" "${MODEL_DSCR_END_LAYER}"
          ;;
      esac
    done
  done
done

echo "[DONE] POPE all-model run finished."
