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

# ------------ CONFIGURATION ------------
CONFIG="configs/slidechat/stage_2_reducer_attn.py"
TEST_CSV="/data/qingq/PathVLM/baselines/github/SlideChat/dataset/SlideBench-VQA-TCGA.csv"
TUMOR_DIR="/data/qingq/PathVLM/baselines/github/SlideChat/dataset/stage_2_vqa_by_tumor/stage2_vqa_tumor_"
OUTPUT_DIR="/data/qingq/PathVLM/baselines/github/SlideChat/models/outputs"
mkdir -p "${OUTPUT_DIR}"
# ---------------------------------------

for CKPT_DIR in "${OUTPUT_DIR}"/*attn; do
  ckpt_name=$(basename "${CKPT_DIR}")
  last_ckpt_file="${CKPT_DIR}/last_checkpoint"

  if [[ ! -f "${last_ckpt_file}" ]]; then
    echo "⚠️  Skipping ${ckpt_name}: no last_checkpoint found."
    continue
  fi

  # 1) Read the exact .pth path and trim whitespace
  ckpt_path=$(< "${last_ckpt_file}")
  ckpt_path="${ckpt_path//[[:space:]]/}"

  # 2) Extract tumor code from folder name (before first '_')
  tumor_lc="${ckpt_name%%_*}"
  tumor=$(echo "${tumor_lc}" | tr '[:lower:]' '[:upper:]')

  # sanity check: does the JSON exist?
  tumor_json="${TUMOR_DIR}/${tumor}.json"
  if [[ ! -f "${tumor_json}" ]]; then
    echo "⚠️  Tumor JSON not found for '${tumor}' (${tumor_json}), skipping."
    continue
  fi

  echo
  echo "=============================================="
  echo " Testing ${tumor} with checkpoint: ${ckpt_path}"
  echo "=============================================="

  EVAL_LOG="/data/qingq/PathVLM/baselines/github/SlideChat/models/outputs/slidechat_reducer_attn_trained_all_eval2.txt"

  out_csv="${OUTPUT_DIR}/output_${tumor_lc}_${ckpt_name}.csv"

  echo "--- ${ckpt_name} | tumor=${tumor} → ${out_csv} ---" | tee -a "${EVAL_LOG}"
  CUDA_VISIBLE_DEVICES=2 \
   xtuner test "${CONFIG}" \
      --checkpoint "${ckpt_path}" \
      --test_slide_csv "${TEST_CSV}" \
      --test_output_csv "${out_csv}" \
      --local_rank 0 \
      --tumor_type "${tumor}" \
      --eval_output_path "${EVAL_LOG}"
done