#!/bin/bash
# Frees disk space after pretraining. Shows what it would remove; deletes only with --go.
# NEVER touches: the two final models (stage1_final, stage2_final), the token data (bin/, stage2/),
# the held-out sets, the code, or anything outside /fsx/ajay-stu/ar-data.
# Usage:  ./cleanup.sh        (dry run, shows sizes)
#         ./cleanup.sh --go   (actually delete)
set -uo pipefail
D=${D:-/fsx/ajay-stu/ar-data}
GO=${1:-}
KEEP1=ckpt_571999        # final Stage 1 (also kept in stage1_final)
KEEP2=ckpt_619699        # final Stage 2 (also kept in stage2_final)

# refuse to delete while training is running
if squeue -u "${USER:-$(id -un)}" -h -n ar-real 2>/dev/null | grep -q .; then
  echo "STOP: a training job is still running"; exit 1
fi

# build the list: training checkpoints except the two finals, and eval exports
targets=()
for p in "$D"/checkpoints/ckpt_*; do
  b=$(basename "$p")
  [ -d "$p" ] && [ "$b" != "$KEEP1" ] && [ "$b" != "$KEEP2" ] && targets+=("$p")
done
for p in "$D"/eval/*.pt; do [ -f "$p" ] && targets+=("$p"); done

if [ ${#targets[@]} -eq 0 ]; then
  echo "nothing to clean"; df -h /fsx | tail -1; exit 0
fi

# show each item with its real disk use (hard-linked copies count as 0 here)
echo "would remove:"
du -shc --apparent-size "${targets[@]}" 2>/dev/null | tail -1 | sed 's/^/  apparent total: /'
for p in "${targets[@]}"; do printf "  %-8s %s\n" "$(du -sh "$p" 2>/dev/null | cut -f1)" "$p"; done
echo "keeping: $D/stage1_final $D/stage2_final $D/bin $D/stage2 $D/heldout"

if [ "$GO" != "--go" ]; then
  echo; echo "DRY RUN - nothing deleted. Run './cleanup.sh --go' to delete."; exit 0
fi

# safety: the two final models must exist before anything is removed
for f in "$D/stage1_final/model_stage1_final.pt" "$D/stage2_final/model_stage2_final.pt" \
         "$D/stage1_final/$KEEP1/.metadata" "$D/stage2_final/$KEEP2/.metadata"; do
  [ -e "$f" ] || { echo "STOP: $f is missing - nothing deleted"; exit 1; }
done
echo "final models verified"
for p in "${targets[@]}"; do rm -rf "$p" && echo "removed $p"; done
df -h /fsx | tail -1