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

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

PYTHON="/opt/miniforge3/envs/xnet_mamba/bin/python"
DATASET="${DATASET:-BUSI}"
RUN_ALL="${RUN_ALL:-0}"
SEED="${SEED:-42}"
EXTRA_SET_ARGS="${EXTRA_SET_ARGS:-}"

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

config_path() {
  case "$1" in
    "BUSI") echo "configs/segmentation/optimized/sup_busi_opt.yaml" ;;
    "BUS-UCLM") echo "configs/segmentation/optimized/sup_bus_uclm_opt.yaml" ;;
    "BUS-BRA") echo "configs/segmentation/optimized/sup_bus_bra_opt.yaml" ;;
    "BUS_UC") echo "configs/segmentation/optimized/sup_bus_uc_opt.yaml" ;;
    "CCAUI") echo "configs/segmentation/optimized/sup_ccaui_opt.yaml" ;;
    "DDTI") echo "configs/segmentation/optimized/sup_ddti_opt.yaml" ;;
    "TN3K") echo "configs/segmentation/optimized/sup_tn3k_opt.yaml" ;;
    "TG3K") echo "configs/segmentation/optimized/sup_tg3k_opt.yaml" ;;
    "OTU_2d") echo "configs/segmentation/optimized/sup_otu_2d_opt.yaml" ;;
    *) echo "Unsupported dataset: $1" >&2; exit 1 ;;
  esac
}

needs_project_split() {
  case "$1" in
    "BUSI"|"BUS-UCLM"|"BUS-BRA"|"BUS_UC"|"CCAUI"|"DDTI") return 0 ;;
    *) return 1 ;;
  esac
}

prepare_split() {
  local dataset="$1"
  local root
  root="$(dataset_root "$dataset")"
  if needs_project_split "$dataset"; then
    echo "[split] ${dataset}"
    "$PYTHON" scripts/generate_project_split.py --dataset "$dataset" --root "$root" --seed "$SEED"
  fi
}

run_one() {
  local dataset="$1"
  local config
  config="$(config_path "$dataset")"
  prepare_split "$dataset"
  "$PYTHON" - "$config" ${EXTRA_SET_ARGS} <<'PY'
import sys

from lib.utils.config import apply_dotlist_overrides, load_yaml_config

config = sys.argv[1]
overrides = sys.argv[2:]
cfg = apply_dotlist_overrides(load_yaml_config(config), overrides)
print("[config]", config)
print(
    "[effective]",
    f"dataset={cfg['dataset']['dataset_name']}",
    f"root={cfg['dataset']['root']}",
    f"image_size={cfg['dataset']['image_size']}",
    f"batch_size={cfg['train']['batch_size']}",
    f"val_batch_size={cfg['train']['val_batch_size']}",
    f"amp={cfg['train']['amp']}",
    f"lr={cfg['optimizer']['lr']}",
    f"swanlab={cfg['logging']['use_swanlab']}",
    f"swanlab_mode={cfg['logging'].get('swanlab_mode')}",
    f"swanlab_logdir={cfg['logging'].get('swanlab_logdir', 'swanlog')}",
    f"experiment={cfg['logging']['experiment_name']}",
)
PY
  echo "[train] ${dataset} using ${config}"
  "$PYTHON" tools/train.py --config "$config" --set ${EXTRA_SET_ARGS}
}

if [[ "$RUN_ALL" == "1" ]]; then
  for dataset in BUSI BUS-UCLM BUS-BRA BUS_UC CCAUI DDTI TN3K TG3K OTU_2d; do
    run_one "$dataset"
  done
else
  run_one "$DATASET"
fi