#!/bin/bash
# ============================================================
# 5-Fold Cross-Validation — SwinUNet, top-20 bands, cosine LR
# Each fold: honest held-out val → checkpoint selected on generalisation
# ============================================================
set -e
source /data1/nasiruddink/miniconda3/etc/profile.d/conda.sh
conda activate brats
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/data1/nasiruddink/miniconda3/envs/brats/lib:/data1/nasiruddink/miniconda3/envs/brats/lib/python3.10/site-packages/nvidia/cudnn/lib
export TF_CPP_MIN_LOG_LEVEL=2

BANDS="band_25,band_18,band_22,band_23,band_10,band_6,band_24,band_27,band_17,band_11,band_13,band_14,band_4,band_3,band_9,band_26,band_8,band_5,band_15,band_7"
mkdir -p bandselection/logs

train_fold() {
    local gpu=$1 fold=$2
    echo "[GPU $gpu] training fold $fold ..."
    CUDA_VISIBLE_DEVICES=$gpu MODEL_NAME="swin_unet" \
      DATA_DIR="dataset/cv_folds/fold_${fold}" \
      BAND_LIST=$BANDS EPOCHS=250 \
      RUN_NAME="cv_fold${fold}" CKPT_DIR="bandselection/cv_fold${fold}" \
      python3 train.py > bandselection/logs/cv_fold${fold}.log 2>&1
    echo "[GPU $gpu] fold $fold done."
}

# Batch 1: folds 0,1,2 on GPU 0,1,2
train_fold 0 0 &  P0=$!
train_fold 1 1 &  P1=$!
train_fold 2 2 &  P2=$!
echo "Batch 1 running (folds 0,1,2)..."
wait $P0 $P1 $P2
echo "Batch 1 done. 60s cooldown."; sleep 60

# Batch 2: folds 3,4 on GPU 0,1
train_fold 0 3 &  P3=$!
train_fold 1 4 &  P4=$!
echo "Batch 2 running (folds 3,4)..."
wait $P3 $P4

echo ""
echo "ALL 5 FOLDS TRAINED."
echo "Checkpoints: bandselection/cv_fold{0..4}/best_model.weights.h5"