#!/bin/bash
# Inference for ct_binary_coronary_segmentation
#
# Usage:
#   ./scripts/inference.sh                                      # Ensemble on default data
#   ./scripts/inference.sh --dataset-dir /path/to/images        # Ensemble on custom data
#   ./scripts/inference.sh --single --fold 0                    # Single model
#   ./scripts/inference.sh --output-dir /path/to/output         # Custom output
#   ./scripts/inference.sh --gpu 1                              # Use specific GPU

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

ENSEMBLE=true
GPU=0
EXTRA_ARGS=()

while [[ $# -gt 0 ]]; do
    case $1 in
        --single) ENSEMBLE=false; shift ;;
        --dataset-dir) EXTRA_ARGS+=("--dataset_dir" "$2"); shift 2 ;;
        --output-dir) EXTRA_ARGS+=("--output_dir" "$2"); shift 2 ;;
        --fold) EXTRA_ARGS+=("--fold" "$2"); shift 2 ;;
        --gpu) GPU="$2"; shift 2 ;;
        *) EXTRA_ARGS+=("$1"); shift ;;
    esac
done

if $ENSEMBLE; then
    CONFIG="configs/ensemble_inference.yaml"
else
    CONFIG="configs/inference.yaml"
fi

echo "=== Running inference ($(if $ENSEMBLE; then echo 'ensemble'; else echo 'single model'; fi)) on GPU $GPU ==="
CUDA_VISIBLE_DEVICES=$GPU micromamba run -n monai python -m monai.bundle run inference \
    --config_file "$CONFIG" \
    "${EXTRA_ARGS[@]+"${EXTRA_ARGS[@]}"}"