#!/usr/bin/env bash
# SUM Parts - fix mode=val crashing before it starts
#
#   main.py:227  validate_fn(model, val_loader, cfg, num_votes=1, epoch=epoch)
#   UnboundLocalError: local variable 'epoch' referenced before assignment
#
# `epoch` is only bound inside the training loop, so the standalone validation
# path references it before it exists. Bind it to -1 (the same sentinel
# validate() already documents in its signature) right before the call.
#
# Idempotent.
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-VAL-EPOCH' "$MAIN"; then
    echo "already patched"
    exit 0
fi

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

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

old = """            if cfg.mode == 'val':
                best_epoch, best_val = load_checkpoint(model, pretrained_path=cfg.pretrained_path)
                val_miou, val_macc, val_oa, val_ious, val_accs = validate_fn(model, val_loader, cfg, num_votes=1, epoch=epoch)"""

new = """            if cfg.mode == 'val':
                best_epoch, best_val = load_checkpoint(model, pretrained_path=cfg.pretrained_path)
                # SUMPARTS-VAL-EPOCH: `epoch` is only bound inside the training
                # loop below, so mode=val referenced it before assignment and
                # died with UnboundLocalError. -1 is the sentinel validate()
                # already defaults to.
                epoch = best_epoch if best_epoch is not None else -1
                val_miou, val_macc, val_oa, val_ious, val_accs = validate_fn(model, val_loader, cfg, num_votes=1, epoch=epoch)"""

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"