#!/bin/bash
# U-Mamba preprocessing for datasets 2-10 (cvc/d1 already done) into a SEPARATE
# preprocessed_umamba dir (U-Mamba's nnunetv2 2.1.x plans differ from the 2.7.0
# nnunet env). Runs in the umamba env. Sequential + thread-capped so the planner's
# BLAS threads don't oversubscribe the 384-core box. Copies the 3-identical-fold
# splits_final.json into each preprocessed dataset afterward.
set -u
cd /home/wzhang/LSC/Code/NPJ
source /opt/anaconda3/etc/profile.d/conda.sh
conda activate umamba
export nnUNet_raw=/home/wzhang/LSC/Code/NPJ/nnunet_workspace/raw
export nnUNet_preprocessed=/home/wzhang/LSC/Code/NPJ/nnunet_workspace/preprocessed_umamba
export nnUNet_results=/home/wzhang/LSC/Code/NPJ/nnunet_workspace/results_umamba
export OMP_NUM_THREADS=2 MKL_NUM_THREADS=2 OPENBLAS_NUM_THREADS=2 NUMEXPR_NUM_THREADS=2
mkdir -p "$nnUNet_preprocessed" "$nnUNet_results"
LOGD=nnunet_workspace/umamba_preprocess_logs; mkdir -p "$LOGD"

dsname () {
  case $1 in
    2) echo Dataset002_kvasir_seg_official;;
    3) echo Dataset003_fives_official;;
    4) echo Dataset004_refuge2_official;;
    5) echo Dataset005_busi_fold01;;
    6) echo Dataset006_idridd_segmentation_fold01;;
    7) echo Dataset007_acdc_png_official;;
    8) echo Dataset008_pannuke_semantic_fold01;;
    9) echo Dataset009_medsegdb_isic2018_holdout;;
    10) echo Dataset010_medsegdb_kits19_fold01;;
  esac
}

for id in 2 3 4 5 6 7 8 9 10; do
  name=$(dsname "$id")
  echo "=== preprocess d$id ($name) ==="
  nnUNetv2_plan_and_preprocess -d "$id" -c 2d -np 16 > "$LOGD/d${id}.log" 2>&1
  rc=$?
  if [ -d "$nnUNet_preprocessed/$name" ]; then
    cp "nnunet_workspace/raw/$name/splits_final.json" "$nnUNet_preprocessed/$name/splits_final.json"
    echo "d$id rc=$rc splits_copied=$( [ -f "$nnUNet_preprocessed/$name/splits_final.json" ] && echo Y || echo N )"
  else
    echo "d$id rc=$rc FAILED_no_preprocessed_dir"
  fi
done
echo UMAMBA_PREPROCESS_DONE