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

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

# ===== 可直接改这里 =====
DATASET="${DATASET:-BUSI}"          # BUS-UCLM | BUSI | BUS-BRA | BUS_UC | CCAUI | DDTI | OTU_2d | TN3K | TG3K
SEED="${SEED:-42}"
RUN_ALL_SUP="${RUN_ALL_SUP:-0}"     # 1 表示跑内置所有全监督实验
PYTHON_BIN="${PYTHON_BIN:-python}"
EXTRA_SET_ARGS="${EXTRA_SET_ARGS:-}"

# ===== 数据集根目录 =====
dataset_root() {
  case "$1" in
    "BUS-UCLM") echo "data/BUS-UCLM" ;;
    "BUSI") echo "data/BUSI" ;;
    "BUS-BRA") echo "data/BUS-BRA" ;;
    "BUS_UC") echo "data/BUS_UC" ;;
    "CCAUI") echo "data/CCAUI" ;;
    "DDTI") echo "data/DDTI" ;;
    "OTU_2d") echo "data/OTU_2d" ;;
    "TN3K") echo "data/TN3K" ;;
    "TG3K") echo "data/TG3K" ;;
    *) echo "Unsupported dataset: $1" >&2; exit 1 ;;
  esac
}

# ===== 是否需要项目级 train/val =====
needs_project_split() {
  case "$1" in
    "BUS-UCLM"|"BUSI"|"BUS-BRA"|"BUS_UC"|"CCAUI"|"DDTI") return 0 ;;
    *) return 1 ;;
  esac
}

prepare_project_splits() {
  local dataset="$1"
  local root
  root="$(dataset_root "$dataset")"

  if needs_project_split "$dataset"; then
    echo "[split] generate project split for ${dataset}"
    "$PYTHON_BIN" tmp/generate_project_split.py --dataset "$dataset" --root "$root" --seed "$SEED"
  fi
}

run_supervised() {
  local dataset="$1"
  local root
  root="$(dataset_root "$dataset")"
  prepare_project_splits "$dataset"
  echo "[train] supervised ${dataset}"
  "$PYTHON_BIN" tools/train.py \
    --config configs/segmentation/train_sup_us_template.yaml \
    --set \
      dataset.dataset_name="$dataset" \
      dataset.root="$root" \
      checkpoint.dir="outputs/experiments/supervised/${dataset}" \
      logging.experiment_name="sup_${dataset}" \
      ${EXTRA_SET_ARGS}
}

run_all_supervised_suite() {
  local datasets=(
    "BUS-UCLM"
    "BUSI"
    "BUS-BRA"
    "BUS_UC"
    "CCAUI"
    "DDTI"
    "OTU_2d"
    "TN3K"
    "TG3K"
  )
  for ds in "${datasets[@]}"; do
    run_supervised "$ds"
  done
}

main() {
  if [[ "$RUN_ALL_SUP" == "1" ]]; then
    run_all_supervised_suite
    exit 0
  fi

  run_supervised "$DATASET"
}

main "$@"