#!/usr/bin/env bash
# SUM Parts - dataset download from Hugging Face
#
# PREREQUISITE: the dataset is GATED (gated: auto). Before this works you must,
# once, in a browser:
#   1. open https://huggingface.co/datasets/gwxgrxhyz/SUM-Parts
#   2. log in and accept the CC BY-NC 4.0 terms on the gate form
# Without that step every download returns HTTP 403.
#
# Usage:
#   bash download_data.sh            # demo only (smoke test, small)
#   bash download_data.sh all        # demo + mesh + pcl (large)
set -euo pipefail

CONDA_ROOT="$HOME/miniconda3"
ENV_NAME="sumparts"
REPO_ID="gwxgrxhyz/SUM-Parts"
# cfgs/sumv2_triangle/default.yaml uses data_root: ../../data/... relative to
# PointNeXt_bundle, which resolves to <repo>/data/
DATA_DIR="$HOME/sum-parts/data"
DL_DIR="$DATA_DIR/_archives"

source "$CONDA_ROOT/etc/profile.d/conda.sh"
conda activate "$ENV_NAME"

python -c "import huggingface_hub" 2>/dev/null || pip install --no-cache-dir "huggingface_hub[cli]"

# Reuse the token already cached on the Windows side if WSL has none.
if [ ! -f "$HOME/.cache/huggingface/token" ] && [ -f /mnt/c/Users/"$USER"/.cache/huggingface/token ]; then
    mkdir -p "$HOME/.cache/huggingface"
    cp /mnt/c/Users/"$USER"/.cache/huggingface/token "$HOME/.cache/huggingface/token"
    chmod 600 "$HOME/.cache/huggingface/token"
    echo "copied HF token from Windows profile"
fi

if [ "${1:-demo}" = "all" ]; then
    FILES=(demo.zip mesh.zip pcl.zip)
else
    FILES=(demo.zip)
fi

mkdir -p "$DL_DIR"

for f in "${FILES[@]}"; do
    echo "=== downloading $f ==="
    python - "$REPO_ID" "$f" "$DL_DIR" <<'PY'
import sys
from huggingface_hub import hf_hub_download
repo_id, filename, out_dir = sys.argv[1:4]
p = hf_hub_download(
    repo_id=repo_id,
    filename=filename,
    repo_type="dataset",
    local_dir=out_dir,
)
print("saved:", p)
PY
done

# Extract with python's zipfile rather than `unzip`: the distro has no unzip
# installed and sudo needs a password here, so apt is not an option.
for f in "${FILES[@]}"; do
    echo "=== extracting $f ==="
    python - "$DL_DIR/$f" "$DATA_DIR" <<'PY'
import sys, zipfile
src, dest = sys.argv[1:3]
with zipfile.ZipFile(src) as z:
    names = z.namelist()
    print(f"  {len(names)} entries")
    z.extractall(dest)
print("  ->", dest)
PY
done

# NOTE: no `find ... | head` here. Under `set -euo pipefail`, head closing the
# pipe sends SIGPIPE to find and the script exits 141 -- reported as a failed
# download even though everything extracted fine.
echo "=== resulting layout (depth 2) ==="
find "$DATA_DIR" -maxdepth 2 -not -path '*/_archives/*' -type d | sort

echo
echo "=== ply counts ==="
for d in "$DATA_DIR"/*/; do
    [ "$(basename "$d")" = "_archives" ] && continue
    n=$(find "$d" -name '*.ply' 2>/dev/null | wc -l)
    [ "$n" -gt 0 ] && printf '%-28s %5s ply\n' "$(basename "$d")/" "$n"
done

echo "DATA DONE"