#!/usr/bin/env bash
# SUM Parts - let test() run on the withheld (unlabeled) test split
#
# The 8 test tiles ship with label = -1 for every point. That is deliberate:
# the test set is blind, and the README says to email predictions to the authors
# for scoring. Training and validation labels are normal (0..12).
#
# main.py only checks `if label is not None`, so the -1 placeholder flows into
# ConfusionMatrix.update, where
#     unique_mapping = true * num_classes + pred
# goes negative and torch.bincount rejects it:
#     RuntimeError: bincount only supports 1-d non-negative integral inputs
#
# The fix treats an all-negative label array as "no ground truth": predictions
# are still produced and written, metrics are simply skipped for that tile.
# This changes no scoring behaviour on labeled data.
#
# Idempotent -- re-running detects the patch is already applied.
set -euo pipefail

source "$HOME/miniconda3/etc/profile.d/conda.sh"
conda activate sumparts

REPO="${1:-$HOME/sum-parts/semantic_segmentation/PointNeXt_bundle}"
MAIN="$REPO/examples/segmentation/main.py"

[ -f "$MAIN" ] || { echo "error: $MAIN not found" >&2; exit 1; }

if grep -q 'SUMPARTS-UNLABELED-TEST' "$MAIN"; then
    echo "already patched"
    exit 0
fi

cp -n "$MAIN" "$MAIN.orig" 2>/dev/null || true

python - "$MAIN" <<'PY'
import sys
from pathlib import Path

p = Path(sys.argv[1])
src = p.read_text(encoding="utf-8")

old = """        coord, feat, label, idx_points, voxel_idx, reverse_idx_part, reverse_idx  = load_data(data_path, cfg)
        if label is not None:
            label = torch.from_numpy(label.astype(int).squeeze()).cuda(non_blocking=True)
"""

new = """        coord, feat, label, idx_points, voxel_idx, reverse_idx_part, reverse_idx  = load_data(data_path, cfg)
        # SUMPARTS-UNLABELED-TEST: the shipped test split carries label = -1 for
        # every point (blind test set; predictions are emailed to the authors for
        # scoring). Passing that to ConfusionMatrix makes true*num_classes+pred
        # negative and torch.bincount raises. Treat it as "no ground truth" so
        # predictions are still produced. Labeled splits are unaffected.
        if label is not None and (label < 0).all():
            logging.info(f'  no ground truth in {os.path.basename(data_path)} '
                         f'(all labels are -1) -- predicting without scoring')
            label = None
        if label is not None:
            label = torch.from_numpy(label.astype(int).squeeze()).cuda(non_blocking=True)
"""

if old not in src:
    print("PATTERN NOT FOUND -- main.py differs from what this patch expects", file=sys.stderr)
    raise SystemExit(1)

p.write_text(src.replace(old, new), encoding="utf-8")
print("patched:", p)
PY

python -c "import ast,sys; ast.parse(open(sys.argv[1], encoding='utf-8').read())" "$MAIN" \
    && echo "syntax OK"

echo "PATCH DONE (original kept at $MAIN.orig)"