#!/bin/bash
# Train ct_binary_coronary_segmentation
#
# Usage:
#   ./scripts/train.sh                        # Train fold 0
#   ./scripts/train.sh --fold 2               # Train specific fold
#   ./scripts/train.sh --all-folds            # Train all 5 folds sequentially
#   ./scripts/train.sh --gpu 1                # Use specific GPU

set -euo pipefail
cd "$(dirname "$0")/.."

FOLD=0
ALL_FOLDS=false
GPU=0
EXTRA_ARGS=()

while [[ $# -gt 0 ]]; do
    case $1 in
        --fold) FOLD="$2"; shift 2 ;;
        --all-folds) ALL_FOLDS=true; shift ;;
        --gpu) GPU="$2"; shift 2 ;;
        *) EXTRA_ARGS+=("$1"); shift ;;
    esac
done

run_fold() {
    local fold=$1
    echo "=== Training fold $fold on GPU $GPU ==="
    CUDA_VISIBLE_DEVICES=$GPU micromamba run -n monai python -m monai.bundle run training \
        --config_file configs/train.yaml \
        --cv_fold "$fold" \
        "${EXTRA_ARGS[@]+"${EXTRA_ARGS[@]}"}"
}

if $ALL_FOLDS; then
    for fold in 0 1 2 3 4; do
        run_fold "$fold"
    done
else
    run_fold "$FOLD"
fi