#!/bin/bash

BASE_URL="https://huggingface.co/datasets/pdearena/Maxwell-3D/resolve/main"
SAVE_DIR="raw"

# Create the raw directory if it doesn't exist
mkdir -p "$SAVE_DIR"

# Function to download a file
download_file() {
    local partition=$1
    local id=$2
    local suffix=$3
    local filename="Maxwell3D_${partition}_${id}${suffix}.h5"
    local url="${BASE_URL}/${filename}?download=true"
    wget -O "${SAVE_DIR}/${filename}" "$url"
}

# Download train files
echo "Downloading train files..."
train_ids=(19124 61748 61991 97473 126288 142103 144307 146242 153264 190639 198010 234461 238834 296423 304936 337393 350040 354348 373330 391455 392020 403221 415374 429773 431752 433143 438911 445636 452245 453445 471651 496019 506325 509578 516076 540888 585598 588419 597635 597861 600790 646044 666597 697145 766153 772245 788727 792179 795128 795933 799532 806439 854274 857042 857466 874947 891593 901448 944048 946191 950647 954413 975533 991670)
for id in "${train_ids[@]}"; do
    download_file "train" "$id" "_100"
done

# Download valid files
echo "Downloading valid files..."
valid_ids=(19224 61849 62091 97573 126388 142203 144407 146342 153364 190739 198110 234561 238934 296523 305036 337493 350140 354448 373430 391555 392120 403321 415473 429873 431852 433243 439011 445736 452345 453545 471751 496119 506425 509678 516176 540988 585698 588519 597735 597961 600890 646144 666697 697245 766253 772345 788827 792278 795228 796033 799632 806539 854374 857142 857566 875047 891693 901548 944148 946291 950747 954513 975633 991770)
for id in "${valid_ids[@]}"; do
    download_file "valid" "$id" ""
done

# Download test files
echo "Downloading test files..."
test_ids=(19324 61948 62191 97673 126488 142303 144507 146442 153464 190839 198210 234661 239034 296623 305136 337593 350240 354548 373530 391655 392220 403421 415573 429973 431952 433343 439111 445836 452445 453645 471851 496219 506525 509778 516276 541088 585798 588619 597835 598061 600990 646244 666797 697345 766353 772445 788927 792378 795328 796133 799732 806639 854474 857242 857666 875147 891793 901648 944248 946391 950847 954613 975733 991870)
for id in "${test_ids[@]}"; do
    download_file "test" "$id" ""
done

echo "Download complete!"