#!/usr/bin/env bash
# SUM Parts - measure what masking absent classes does to the Seosan prediction
#
# Runs inference twice on the same tile and the same checkpoint: once as-is,
# once with water and boat masked out of the argmax. The interesting number is
# where those 26% of points land when they can no longer be called water.
set -uo pipefail

SCRIPTS="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
OUT="$HOME/sum-parts/runs/mask_experiment"
mkdir -p "$OUT"

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

# 4 = water, 6 = boat  (openpoints/dataset/sumv2_triangle/sumv2_triangle.py)
MASK="${MASK:-4,6}"

run_and_dump() {
    local tag="$1" maskval="$2"
    echo
    echo "════ $tag ════"
    SUMPARTS_MASK_CLASSES="$maskval" bash "$SCRIPTS/poc_infer.sh" \
        > "$OUT/infer_${tag}.log" 2>&1
    local rc=$?
    if [ $rc -ne 0 ]; then
        echo "  inference FAILED rc=$rc"
        tail -12 "$OUT/infer_${tag}.log"
        return 1
    fi
    grep -a 'masking classes' "$OUT/infer_${tag}.log" | head -1
    bash "$SCRIPTS/dump_seosan_pred.sh" 2>&1 | tee "$OUT/dump_${tag}.txt" \
        | sed -n '/label field/,/^$/p'
}

run_and_dump baseline ""
run_and_dump masked   "$MASK"

echo
echo "════ comparison ════"
python - "$OUT/dump_baseline.txt" "$OUT/dump_masked.txt" <<'PY'
import re
import sys
from pathlib import Path

def parse(path):
    out = {}
    for line in Path(path).read_text(encoding="utf-8", errors="replace").splitlines():
        m = re.match(r"\s+(\d+)\s+(\S+)\s+([\d,]+)\s+([\d.]+)%", line)
        if m:
            out[m.group(2)] = (int(m.group(3).replace(",", "")), float(m.group(4)))
    return out

a, b = parse(sys.argv[1]), parse(sys.argv[2])
names = sorted(set(a) | set(b), key=lambda n: -max(a.get(n, (0,))[0], b.get(n, (0,))[0]))

print(f"  {'class':<20} {'baseline':>12} {'masked':>12} {'change':>12}")
for n in names:
    an, ap = a.get(n, (0, 0.0))
    bn, bp = b.get(n, (0, 0.0))
    d = bn - an
    arrow = "" if d == 0 else ("  <- absorbed" if d > 0 else "")
    print(f"  {n:<20} {ap:>10.2f}% {bp:>10.2f}% {d:>+11,}{arrow}")
PY

echo
echo "logs in $OUT"