#!/usr/bin/env bash
# SUM Parts - find the largest voxel_max that still fits in VRAM
#
# Why this matters: on WSL2 the NVIDIA driver spills past VRAM into host RAM
# instead of raising OOM, so an oversized config still "works" -- at 47 s/iter
# instead of 0.43. Dropping voxel_max until the peak fits gave a 108x speedup
# on pointnext-xl. But voxel_max is also the cfg value the paper trained with
# (64000), so anything lower is a deviation worth naming. This sweep finds how
# close to 64000 the card can actually get.
set -uo pipefail

CONDA_ROOT="$HOME/miniconda3"
SEG="$HOME/sum-parts/semantic_segmentation/PointNeXt_bundle/examples/segmentation"
SCRIPTS="/mnt/d/MYCLAUDE_PROJECT/sum-parts-test/scripts"

CFG="${CFG:-pointnext-xl}"
ITERS="${ITERS:-5}"
VALUES="${*:-24000 32000 40000 48000 64000}"

source "$CONDA_ROOT/etc/profile.d/conda.sh"
conda activate sumparts
export WANDB_MODE=disabled WANDB_SILENT=true CUDA_HOME="$CONDA_PREFIX"

cd "$SEG"

echo "sweeping $CFG over voxel_max: $VALUES"
echo
printf '%-10s %-12s %-11s %-10s %s\n' "voxel_max" "pts/batch" "peak VRAM" "s/iter" "fits?"
printf -- '---------------------------------------------------------\n'

for vm in $VALUES; do
    log="/tmp/sweep_${CFG}_${vm}.log"
    python -u "$SCRIPTS/bench_models.py" \
        --iters "$ITERS" --voxel-max "$vm" --cfgs "$CFG" > "$log" 2>&1
    rc=$?
    if [ $rc -ne 0 ]; then
        printf '%-10s %s\n' "$vm" "FAILED (rc=$rc, see $log)"
        continue
    fi
    line=$(grep -E "^${CFG} +[0-9]" "$log" | tail -1)
    if [ -z "$line" ]; then
        printf '%-10s %s\n' "$vm" "no result (see $log)"
        continue
    fi
    # cfg params bs pts peak s/iter fits...
    pts=$(echo "$line" | awk '{print $4}')
    peak=$(echo "$line" | awk '{print $5}')
    sec=$(echo "$line" | awk '{print $6}')
    fits=$(echo "$line" | cut -d' ' -f7- | sed 's/^ *//')
    printf '%-10s %-12s %-11s %-10s %s\n' "$vm" "$pts" "$peak" "$sec" "$fits"
done

echo
echo "SWEEP DONE"