#!/usr/bin/env bash
# Unpack all tar archives after downloading from HuggingFace
# Usage: cd dREPA_collections && bash unpack.sh
set -euo pipefail

echo "=== Unpacking dREPA collections ==="

# --- COCO256 train ---
if ls data/coco256_features_sdvae_ft_ema/train*.tar 1>/dev/null 2>&1; then
    echo "[1/3] Unpacking COCO256 train..."
    mkdir -p data/coco256_features_sdvae_ft_ema/train
    for f in data/coco256_features_sdvae_ft_ema/train*.tar; do
        echo "  -> $(basename "$f")"
        tar xf "$f" -C data/coco256_features_sdvae_ft_ema/train/
    done
    echo "  Done. $(ls data/coco256_features_sdvae_ft_ema/train/ | wc -l) files unpacked."
else
    echo "[1/3] COCO256 train: no tar files found, skip."
fi

# --- COCO256 val ---
if ls data/coco256_features_sdvae_ft_ema/val*.tar 1>/dev/null 2>&1; then
    echo "[2/3] Unpacking COCO256 val..."
    mkdir -p data/coco256_features_sdvae_ft_ema/val
    for f in data/coco256_features_sdvae_ft_ema/val*.tar; do
        echo "  -> $(basename "$f")"
        tar xf "$f" -C data/coco256_features_sdvae_ft_ema/val/
    done
    echo "  Done. $(ls data/coco256_features_sdvae_ft_ema/val/ | wc -l) files unpacked."
else
    echo "[2/3] COCO256 val: no tar files found, skip."
fi

# --- Metrics data ---
if [ -f metrics/data.tar ]; then
    echo "[3/3] Unpacking metrics data..."
    mkdir -p metrics/data
    tar xf metrics/data.tar -C metrics/data/
    echo "  Done."
else
    echo "[3/3] Metrics data: no tar file found, skip."
fi

echo ""
echo "=== All done! ==="
echo ""
read -p "Remove tar files to save space? [y/N] " -n 1 -r
echo
if [[ $REPLY =~ ^[Yy]$ ]]; then
    rm -f data/coco256_features_sdvae_ft_ema/train*.tar
    rm -f data/coco256_features_sdvae_ft_ema/val*.tar
    rm -f metrics/data.tar
    echo "Tar files removed."
fi