#!/bin/bash
# U-Mamba PanNuke fold02(d11)+fold03(d12): train fold0 (100ep) on both A100s, predict, score.
set -u
cd /home/wzhang/LSC/Code/NPJ
source /opt/anaconda3/etc/profile.d/conda.sh
RAW=/home/wzhang/LSC/Code/NPJ/nnunet_workspace/raw
PRED=/home/wzhang/LSC/Code/NPJ/nnunet_workspace/predTs_umamba
DATA_ROOT=/home/wzhang/LSC/Dataset/Segmentation/processed_unified
TR=nnUNetTrainerUMambaBot_100epochs
export CUDA_DEVICE_ORDER=PCI_BUS_ID
mkdir -p "$PRED"

conda activate umamba
export nnUNet_raw=$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 nnUNet_n_proc_DA=8 OMP_NUM_THREADS=4
CUDA_VISIBLE_DEVICES=4 nnUNetv2_train 11 2d 0 -tr "$TR" > nnunet_workspace/train_umamba_d11.log 2>&1 &
CUDA_VISIBLE_DEVICES=5 nnUNetv2_train 12 2d 0 -tr "$TR" > nnunet_workspace/train_umamba_d12.log 2>&1 &
wait
echo UMAMBA_CV_TRAIN_DONE
CUDA_VISIBLE_DEVICES=4 nnUNetv2_predict -i "$RAW/Dataset011_pannuke_semantic_fold02/imagesTs" -o "$PRED/d11_f0" -d 11 -c 2d -f 0 -tr "$TR" --disable_tta > nnunet_workspace/pred_umamba_d11.log 2>&1 &
CUDA_VISIBLE_DEVICES=5 nnUNetv2_predict -i "$RAW/Dataset012_pannuke_semantic_fold03/imagesTs" -o "$PRED/d12_f0" -d 12 -c 2d -f 0 -tr "$TR" --disable_tta > nnunet_workspace/pred_umamba_d12.log 2>&1 &
wait

conda deactivate; conda activate seggen
export OMP_NUM_THREADS=8 MKL_NUM_THREADS=8 OPENBLAS_NUM_THREADS=8
python framework/nnunet_eval.py --data_root "$DATA_ROOT" --dataset pannuke_semantic --protocol fold02 --raw "$RAW" --dataset_id 11 --fold 0 --pred_dir "$PRED/d11_f0" --arch umamba --exp_name baselines
python framework/nnunet_eval.py --data_root "$DATA_ROOT" --dataset pannuke_semantic --protocol fold03 --raw "$RAW" --dataset_id 12 --fold 0 --pred_dir "$PRED/d12_f0" --arch umamba --exp_name baselines
echo UMAMBA_CV_DONE