#!/usr/bin/env bash
# Pull everything run_longeval.sh needs from HuggingFace onto a fresh Hopper box.
# Checkpoints are ~3.9 GB each (~16 GB total, ~6 min each on a typical link).
set -uo pipefail
cd "$(dirname "$0")"
HFR=${HFR:-unparallelled/BluScriptCP}
HF=${HF:-hf}
CK=checkpoints_bluscriptcp/rank_0
mkdir -p "$CK" Data_Loader/Data

command -v "$HF" >/dev/null || { echo "ERROR: 'hf' CLI not on PATH (pip install -U huggingface_hub)"; exit 1; }
"$HF" auth whoami >/dev/null 2>&1 || { echo "ERROR: run '$HF auth login' first"; exit 1; }

echo "=== checkpoints ==="
declare -A MAP=(
  [extended_blumodelcp_run7_step_750.ckpt]=blumodelcp_run7_step_750.ckpt
  [arm_a_blumodelcp_run12_step_400.ckpt]=blumodelcp_run12_step_400.ckpt
  [arm_b_blumodelcp_run15_step_400.ckpt]=blumodelcp_run15_step_400.ckpt
  [arm_c_blumodelcp_run16_step_400.ckpt]=blumodelcp_run16_step_400.ckpt
)
for SRC in "${!MAP[@]}"; do
  DST="$CK/${MAP[$SRC]}"
  if [ -f "$DST" ]; then echo "  have ${MAP[$SRC]}"; continue; fi
  echo "  downloading $SRC"
  "$HF" download "$HFR" "checkpoints/in2/$SRC" --local-dir /tmp/hf_ck >/dev/null \
    && cp "/tmp/hf_ck/checkpoints/in2/$SRC" "$DST" && echo "    -> $DST" \
    || echo "    FAILED $SRC"
done

# Haystack filler. The generator falls back to synthetic text without it, but real
# flux text is the distribution the model was trained on, so prefer it.
echo "=== haystack (flux val shard) ==="
if ls Data_Loader/Data/flux_val_*.bin >/dev/null 2>&1; then
  echo "  have $(ls Data_Loader/Data/flux_val_*.bin | head -1)"
else
  "$HF" download navingv/flux --repo-type dataset --include "flux_val_000000.bin" \
      --local-dir /tmp/flux >/dev/null 2>&1 \
    && cp /tmp/flux/flux_val_000000.bin Data_Loader/Data/ && echo "  -> Data_Loader/Data/flux_val_000000.bin" \
    || echo "  WARN: no flux shard; generator will use synthetic filler"
fi

echo "=== python deps ==="
python3 -c "import tiktoken" 2>/dev/null && echo "  tiktoken ok" || pip install -q tiktoken && echo "  tiktoken installed"

echo
echo "Next:"
echo "  make CP_FUSED_ROPE=1 CP_ATTN_SPLIT=1 SM_ARCH=90 bluscript-cp -j\$(nproc)"
echo "  ./run_longeval.sh"