#!/bin/bash
# Downloads the pancreas segmentation datasets used by this project:
#   - MSD Task07 Pancreas (CT, Medical Segmentation Decathlon)
#   - AMOS22 (CT + MRI, 15-organ segmentation)
#   - PanSegData (pancreas MRI, T1 + T2, from the PanSegNet study)
#
# Total download size is roughly 40 GiB. Each archive is downloaded to a
# scratch directory, extracted into its own folder here, then deleted.
set -e

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
TMP_DIR="$SCRIPT_DIR/.tmp_downloads"

# Downloads $url, extracts it (tar or zip) into $SCRIPT_DIR/$target_dir, and
# leaves a marker file behind so re-running the script skips finished work.
download_and_extract() {
  local url="$1"
  local archive_name="$2"
  local target_dir="$3"
  local archive_path="$TMP_DIR/$archive_name"
  local marker="$SCRIPT_DIR/$target_dir/.extracted_$archive_name"

  if [ -f "$marker" ]; then
    echo "Skipping $archive_name (already extracted)"
    return
  fi

  mkdir -p "$TMP_DIR" "$SCRIPT_DIR/$target_dir"
  curl -L --fail -C - -o "$archive_path" "$url"

  case "$archive_name" in
    *.tar) tar -xf "$archive_path" -C "$SCRIPT_DIR/$target_dir" ;;
    *.zip) unzip -o "$archive_path" -d "$SCRIPT_DIR/$target_dir" ;;
    *) echo "Don't know how to extract $archive_name" >&2; exit 1 ;;
  esac

  rm "$archive_path"
  touch "$marker"
}

# Downloads a standalone file (e.g. metadata) with no extraction step.
download_file() {
  local url="$1"
  local file_name="$2"
  local target_dir="$3"
  local dest="$SCRIPT_DIR/$target_dir/$file_name"

  if [ -f "$dest" ]; then
    echo "Skipping $file_name (already downloaded)"
    return
  fi

  mkdir -p "$SCRIPT_DIR/$target_dir"
  curl -L --fail -o "$dest" "$url"
}

# MSD Task07 Pancreas (~11.4 GiB) — CT scans with pancreas outlines
download_and_extract \
  "https://msd-for-monai.s3-us-west-2.amazonaws.com/Task07_Pancreas.tar" \
  "Task07_Pancreas.tar" \
  "MSD_Task07_Pancreas"

# AMOS22 (~22.6 GiB) — CT + MRI, 15-organ segmentation (incl. pancreas)
download_and_extract \
  "https://zenodo.org/records/7262581/files/amos22.zip?download=1" \
  "amos22.zip" \
  "AMOS22"
download_file \
  "https://zenodo.org/records/7262581/files/labeled_data_meta_0000_0599.csv?download=1" \
  "labeled_data_meta_0000_0599.csv" \
  "AMOS22"

# PanSegData (T1 ~4.0 GiB, T2 ~1.4 GiB) — pancreas MRI, PanSegNet study
download_and_extract \
  "https://osf.io/download/ch8ay/" \
  "t1.zip" \
  "PanSegData/T1"
download_and_extract \
  "https://osf.io/download/bre8p/" \
  "t2.zip" \
  "PanSegData/T2"

rmdir "$TMP_DIR" 2>/dev/null || true

echo "All downloads complete."
