#!/usr/bin/env bash
# nnU-Net + U-Mamba @512 re-eval on a100 (CPU-only, does NOT touch GPUs).
# Re-scores cached predTs/predTs_umamba at eval_size 512 into results/unified512/
# so they join the unified-512 table. ~64 cells, parallel (concurrency MAXJ).
set -u
cd /home/wzhang/LSC/Code/NPJ
source /opt/anaconda3/etc/profile.d/conda.sh
conda activate seggen
export OMP_NUM_THREADS=4 MKL_NUM_THREADS=4 OPENBLAS_NUM_THREADS=4 NUMEXPR_NUM_THREADS=4
DR=/home/wzhang/LSC/Dataset/Segmentation/processed_unified
RAW=/home/wzhang/LSC/Code/NPJ/nnunet_workspace/raw
PRED_NN=/home/wzhang/LSC/Code/NPJ/nnunet_workspace/predTs
PRED_UM=/home/wzhang/LSC/Code/NPJ/nnunet_workspace/predTs_umamba
MAXJ=10

declare -A DS=(
  [1]=cvc_clinicdb:official [2]=kvasir_seg:official [3]=fives:official
  [4]=refuge2:official [5]=busi:fold01 [6]=idridd_segmentation:fold01
  [7]=acdc_png:official [8]=pannuke_semantic:fold01 [9]=medsegdb_isic2018:holdout
  [10]=medsegdb_kits19:fold01 [11]=pannuke_semantic:fold02 [12]=pannuke_semantic:fold03
)

run=0
for id in 1 2 3 4 5 6 7 8 9 10 11 12; do
  IFS=: read -r ds proto <<< "${DS[$id]}"
  for f in 0 1 2; do
    for ap in "nnunet:$PRED_NN" "umamba:$PRED_UM"; do
      arch=${ap%%:*}; pred=${ap#*:}
      outdir=$pred/d${id}_f${f}
      [ -d "$outdir" ] && ls -A "$outdir"/*.png >/dev/null 2>&1 || continue
      (
        python framework/nnunet_eval.py --data_root "$DR" --dataset "$ds" --protocol "$proto" \
          --raw "$RAW" --dataset_id "$id" --fold "$f" --pred_dir "$outdir" --arch "$arch" \
          --exp_name unified512 --eval_size 512 \
          > /tmp/nnum512_${arch}_d${id}_f${f}.log 2>&1 \
          && echo "[ok] $arch d${id}_f${f} ($ds $proto)" \
          || echo "[FAIL] $arch d${id}_f${f} ($ds $proto)"
      ) &
      run=$((run+1))
      if [ "$run" -ge "$MAXJ" ]; then wait -n; run=$((run-1)); fi
    done
  done
done
wait
echo "NNUM_EVAL512_DONE"
n=$(find results/unified512 -path '*/nnunet/*/metrics.json' -o -path '*/umamba/*/metrics.json' 2>/dev/null | wc -l)
echo "unified512 nnunet+umamba metrics.json count: $n"