#!/usr/bin/env bash
# build_robustbench_subset.sh — construct the official RobustBench 5K subset
# of ImageNet-1k validation as symlinks.
#
# Prerequisites:
#   1. ImageNet 2012 validation set (50K images, ~6.7 GB) downloaded and
#      extracted to a directory with structure:
#          <imagenet_val>/n01440764/ILSVRC2012_val_00000293.JPEG
#          <imagenet_val>/n01440764/ILSVRC2012_val_00002138.JPEG
#          ...
#      (one subdir per synset, ImageFolder convention)
#
#   2. The RobustBench repo cloned somewhere (we just need one file from it).
#
# Usage:
#   bash scripts/build_robustbench_subset.sh <imagenet_val_dir> [out_dir]
#
# Example:
#   bash scripts/build_robustbench_subset.sh /Volumes/Datasets/imagenet_val data/robustbench_subset
#
# Output:
#   <out_dir>/n09468604/ILSVRC2012_val_00015253.JPEG  (symlink → original)
#   ... (4999 symlinks total, paper rounds to 5000)
#
# Why symlinks: avoids duplicating ~7GB of data; the RobustBench loader and
# our run_attack_sweep.py both work with ImageFolder structure regardless.

set -euo pipefail

if [ "$#" -lt 1 ]; then
  echo "Usage: $0 <imagenet_val_dir> [out_dir]"
  echo ""
  echo "  <imagenet_val_dir>   path to ImageNet 2012 val (50K, ImageFolder layout)"
  echo "  [out_dir]            output dir (default: data/robustbench_subset)"
  exit 1
fi

IMAGENET_VAL="$(realpath "$1")"
OUT_DIR="${2:-data/robustbench_subset}"
PROJECT_ROOT="$(cd "$(dirname "$0")/.." && pwd)"
OUT_DIR="$PROJECT_ROOT/$OUT_DIR"

# Sanity: source must exist and look like ImageNet val
if [ ! -d "$IMAGENET_VAL" ]; then
  echo "ERROR: $IMAGENET_VAL does not exist"
  exit 1
fi
N_SYNSETS=$(find "$IMAGENET_VAL" -maxdepth 1 -mindepth 1 -type d | wc -l | tr -d ' ')
if [ "$N_SYNSETS" -lt 100 ]; then
  echo "ERROR: $IMAGENET_VAL contains $N_SYNSETS synset dirs — expected 1000"
  echo "  Make sure the structure is <imagenet_val>/<synset>/<file>.JPEG"
  exit 1
fi
echo "Source ImageNet val: $IMAGENET_VAL ($N_SYNSETS synset dirs)"

# Get the IDs file from RobustBench repo (clone temp if not provided)
TMP_REPO="/tmp/robustbench_for_subset"
if [ ! -d "$TMP_REPO" ]; then
  echo "Cloning RobustBench (depth=1) to /tmp ..."
  git clone --depth 1 https://github.com/RobustBench/robustbench "$TMP_REPO" 2>&1 | tail -3
fi
IDS_FILE="$TMP_REPO/robustbench/helper_files/imagenet_test_image_ids.txt"
if [ ! -f "$IDS_FILE" ]; then
  echo "ERROR: $IDS_FILE not found"
  exit 1
fi
N_IDS=$(wc -l < "$IDS_FILE" | tr -d ' ')
echo "RobustBench IDs file: $IDS_FILE ($N_IDS lines)"

# Build the subset
mkdir -p "$OUT_DIR"
echo "Building subset at: $OUT_DIR"

linked=0
missing=0
while IFS= read -r line; do
  # Skip empty lines
  [ -z "$line" ] && continue

  src="$IMAGENET_VAL/$line"
  dst="$OUT_DIR/$line"

  if [ ! -f "$src" ]; then
    missing=$((missing + 1))
    if [ $missing -le 3 ]; then
      echo "  WARNING: missing source: $src"
    fi
    continue
  fi

  mkdir -p "$(dirname "$dst")"
  ln -sf "$src" "$dst"
  linked=$((linked + 1))
done < "$IDS_FILE"

echo ""
echo "Done."
echo "  Linked:  $linked / $N_IDS"
echo "  Missing: $missing"

if [ $missing -gt 0 ]; then
  echo "  WARNING: $missing source files were not found in $IMAGENET_VAL"
  echo "  Check that the ImageNet val download is complete and uses synset/file structure."
fi

# Verify
final_count=$(find "$OUT_DIR" -name "*.JPEG" | wc -l | tr -d ' ')
echo "  Files in subset: $final_count (expected ~4999)"