#!/bin/bash
# U-Mamba finale: predict test sets (umamba env, both A100s) then score the 7
# framework metrics (seggen env, thread-capped). Mirrors the nnU-Net predict+eval
# flow but with results_umamba / preprocessed_umamba / trainer 100ep / arch=umamba.
set -u
cd /home/wzhang/LSC/Code/NPJ
source /opt/anaconda3/etc/profile.d/conda.sh
export CUDA_DEVICE_ORDER=PCI_BUS_ID
RAW=/home/wzhang/LSC/Code/NPJ/nnunet_workspace/raw
PRE=/home/wzhang/LSC/Code/NPJ/nnunet_workspace/preprocessed_umamba
RES=/home/wzhang/LSC/Code/NPJ/nnunet_workspace/results_umamba
PRED=/home/wzhang/LSC/Code/NPJ/nnunet_workspace/predTs_umamba
DATA_ROOT=/home/wzhang/LSC/Dataset/Segmentation/processed_unified
TR=nnUNetTrainerUMambaBot_100epochs
mkdir -p "$PRED"

dsname () { case $1 in
  1) echo cvc_clinicdb official;;     2) echo kvasir_seg official;;  3) echo fives official;;
  4) echo refuge2 official;;          5) echo busi fold01;;          6) echo idridd_segmentation fold01;;
  7) echo acdc_png official;;         8) echo pannuke_semantic fold01;; 9) echo medsegdb_isic2018 holdout;;
  10) echo medsegdb_kits19 fold01;;
esac; }
rawname () { printf "Dataset%03d_%s_%s" "$1" "$2" "$3"; }

# ---------- Phase 1: predict (umamba env), ids split across GPU4 / GPU5 ----------
conda activate umamba
export nnUNet_raw=$RAW nnUNet_preprocessed=$PRE nnUNet_results=$RES
export OMP_NUM_THREADS=4 MKL_NUM_THREADS=4
predict_ids () {
  local gpu=$1; shift
  for id in "$@"; do
    read -r ds proto < <(dsname "$id"); dn=$(rawname "$id" "$ds" "$proto")
    for f in 0 1 2; do
      out=$PRED/d${id}_f${f}; mkdir -p "$out"
      echo "[predict gpu$gpu] $dn f$f"
      CUDA_VISIBLE_DEVICES=$gpu nnUNetv2_predict -i "$RAW/$dn/imagesTs" -o "$out" \
        -d "$id" -c 2d -f "$f" -tr "$TR" --disable_tta > "$out/predict.log" 2>&1
    done
  done
}
predict_ids 4 1 2 3 4 5 &
predict_ids 5 6 7 8 9 10 &
wait
echo PREDICT_DONE

# ---------- Phase 2: eval (seggen env), parallel, thread-capped, max 10 ----------
conda deactivate; conda activate seggen
export OMP_NUM_THREADS=8 MKL_NUM_THREADS=8 OPENBLAS_NUM_THREADS=8 NUMEXPR_NUM_THREADS=8
for id in 1 2 3 4 5 6 7 8 9 10; do
  read -r ds proto < <(dsname "$id")
  for f in 0 1 2; do
    while (( $(jobs -rp | wc -l) >= 10 )); do wait -n; done
    (
      python framework/nnunet_eval.py --data_root "$DATA_ROOT" \
        --dataset "$ds" --protocol "$proto" --raw "$RAW" \
        --dataset_id "$id" --fold "$f" --pred_dir "$PRED/d${id}_f${f}" \
        --arch umamba --exp_name baselines > "$PRED/d${id}_f${f}/eval.log" 2>&1
      echo "evaled d${id}_f${f}: $(tail -1 "$PRED/d${id}_f${f}/eval.log")"
    ) &
  done
done
wait
echo EVAL_DONE