#!/usr/bin/env bash
# SUM Parts - stop ConfusionMatrix.update from overwriting its caller's predictions
#
# openpoints/utils/metrics.py:
#
#     if (true == self.ignore_index).sum() > 0:
#         pred[true == self.ignore_index] = self.virtual_num_classes - 1
#         true[true == self.ignore_index] = self.virtual_num_classes - 1
#
# Folding ignored points into the last bucket is fine for scoring, but it is
# done in place on the tensors the caller passed in. main.py's test() calls
# cm.update(pred, label) and *then* writes the visualization from that same
# pred, so the file on disk is the mutated copy, not what the model predicted.
#
# On a tile labelled entirely with the ignore class - which is exactly what an
# unlabelled tile of our own data looks like, label=0 with ignore_index=0 -
# every single point gets rewritten to class num_classes-1 (wall). The
# prediction file then reads 100% wall no matter what the network actually
# said, and the model looks degenerate when it is not.
#
# Fix: clone before masking. Scoring is unchanged; the caller's tensors survive.
#
# Idempotent.
set -euo pipefail

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

REPO="${1:-$HOME/sum-parts/semantic_segmentation/PointNeXt_bundle}"
METRICS="$REPO/openpoints/utils/metrics.py"

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

if grep -q 'SUMPARTS-NO-MUTATE' "$METRICS"; then
    echo "already patched"
    exit 0
fi

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

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

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

old = """        true = true.flatten()
        pred = pred.flatten()
        if self.ignore_index is not None:
            if (true == self.ignore_index).sum() > 0:
                pred[true == self.ignore_index] = self.virtual_num_classes -1
                true[true == self.ignore_index] = self.virtual_num_classes -1"""

new = """        # SUMPARTS-NO-MUTATE: clone before masking. These used to be written
        # in place, which silently rewrote the caller's prediction tensor --
        # main.py's test() saves its visualization from the same `pred` right
        # after calling this, so the file on disk showed the folded values
        # rather than the model's output. On a tile labelled entirely with the
        # ignore class every point came out as num_classes-1.
        true = true.flatten().clone()
        pred = pred.flatten().clone()
        if self.ignore_index is not None:
            if (true == self.ignore_index).sum() > 0:
                pred[true == self.ignore_index] = self.virtual_num_classes -1
                true[true == self.ignore_index] = self.virtual_num_classes -1"""

if old not in src:
    print("PATTERN NOT FOUND -- metrics.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())" "$METRICS" \
    && echo "syntax OK"

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