#!/usr/bin/env python3
"""Run prompt-consensus accessibility masks on a small reviewed pilot set.

Supported backends:
  sam3          Meta SAM 3 text-prompt concept segmentation (preferred)
  grounded_sam  GroundingDINO boxes refined by SAM ViT-H (local control)

The tool saves per-prompt masks, a high-recall union, a majority-vote core, an
uncertainty mask, both raw/high-recall and filtered obstacle masks, overlays,
and machine-readable quality scores.  Automatic output is a review proposal,
never ground truth.
"""

from __future__ import annotations

import argparse
import html
import json
import math
import os
import re
import sys
import traceback
from collections import Counter
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import cv2
import numpy as np
from PIL import Image, ImageOps


def read_jsonl(path: Path) -> list[dict[str, Any]]:
    rows: list[dict[str, Any]] = []
    with path.open("r", encoding="utf-8") as handle:
        for line in handle:
            if line.strip():
                rows.append(json.loads(line))
    return rows


def write_jsonl(path: Path, rows: list[dict[str, Any]]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    tmp_path = path.with_suffix(path.suffix + ".tmp")
    with tmp_path.open("w", encoding="utf-8") as handle:
        for row in rows:
            handle.write(json.dumps(row, ensure_ascii=False, sort_keys=True) + "\n")
    tmp_path.replace(path)


def json_dump(path: Path, value: Any) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    tmp_path = path.with_suffix(path.suffix + ".tmp")
    with tmp_path.open("w", encoding="utf-8") as handle:
        json.dump(value, handle, ensure_ascii=False, indent=2, sort_keys=True)
    tmp_path.replace(path)


STANDARD_SAMPLE_OUTPUTS = (
    "target_visible_candidate.png",
    "target_visible_core.png",
    "target_uncertain.png",
    "obstacle_unfiltered_evidence.png",
    "obstacle_all_detected.png",
    "obstacle.png",
    "obstacle_red.png",
    "overlay.png",
    "quality.json",
)
OBSTACLE_EVIDENCE_POLICY_VERSION = 1


def load_resume_quality(sample_dir: Path, sample_id: str, backend_name: str) -> dict[str, Any] | None:
    quality_path = sample_dir / "quality.json"
    if not quality_path.is_file():
        return None
    if any(not (sample_dir / filename).is_file() for filename in STANDARD_SAMPLE_OUTPUTS):
        return None
    with quality_path.open("r", encoding="utf-8") as handle:
        quality = json.load(handle)
    if quality.get("sample_id") != sample_id:
        return None
    if quality.get("backend") != backend_name:
        return None
    # Old output roots did not distinguish raw prompt evidence from the
    # high-recall-but-sanity-filtered evidence used by later hidden completion.
    # Force an explicit rerun rather than silently reusing that weaker audit
    # trail after this schema change.
    if quality.get("obstacle_evidence_policy_version") != OBSTACLE_EVIDENCE_POLICY_VERSION:
        return None
    if quality.get("review_status") == "error" or quality.get("error"):
        return None
    return quality


def safe_name(text: str) -> str:
    return re.sub(r"[^a-z0-9]+", "_", text.lower()).strip("_")


def prompt_list(value: Any) -> list[str]:
    """Normalize a prompt list while accepting the legacy JSON shape."""
    if isinstance(value, str):
        raw = [value]
    elif isinstance(value, (list, tuple)):
        raw = value
    else:
        raw = []
    return list(dict.fromkeys(str(prompt).strip() for prompt in raw if str(prompt).strip()))


def obstacle_prompt_resolution(
    config: dict[str, Any], category: str, backend_name: str | None = None
) -> dict[str, Any]:
    """Resolve backend-aware obstacle prompts without breaking legacy configs.

    When a backend-specific list is present it replaces the corresponding
    legacy global/category list.  The other level still falls back separately,
    so a config may specialize only global prompts or only one category.  With
    ``backend_name=None`` the output is exactly the historical legacy union.
    """
    global_prompts = prompt_list(config.get("obstacle_prompts"))
    category_map = config.get("category_obstacle_prompts")
    category_prompts = prompt_list(
        category_map.get(category) if isinstance(category_map, dict) else None
    )
    global_source = "legacy.obstacle_prompts"
    category_source = "legacy.category_obstacle_prompts"
    if backend_name:
        by_backend = config.get("obstacle_prompts_by_backend")
        if isinstance(by_backend, dict) and backend_name in by_backend:
            global_prompts = prompt_list(by_backend.get(backend_name))
            global_source = f"obstacle_prompts_by_backend.{backend_name}"
        category_by_backend = config.get("category_obstacle_prompts_by_backend")
        backend_categories = (
            category_by_backend.get(backend_name)
            if isinstance(category_by_backend, dict)
            else None
        )
        if isinstance(backend_categories, dict) and category in backend_categories:
            category_prompts = prompt_list(backend_categories.get(category))
            category_source = (
                f"category_obstacle_prompts_by_backend.{backend_name}.{category}"
            )
    prompts = list(dict.fromkeys([*global_prompts, *category_prompts]))
    return {
        "backend": backend_name,
        "category": category,
        "global_source": global_source,
        "category_source": category_source,
        "prompts": prompts,
    }


def obstacle_prompts_for_category(
    config: dict[str, Any], category: str, backend_name: str | None = None
) -> list[str]:
    """Return backend-aware obstacle prompts, with legacy fallback support."""
    return list(obstacle_prompt_resolution(config, category, backend_name)["prompts"])


def as_mapping(value: Any) -> dict[str, Any]:
    return dict(value) if isinstance(value, dict) else {}


def merge_mappings(base: dict[str, Any], override: dict[str, Any]) -> dict[str, Any]:
    """Small recursive merge for backend-specific evidence policy overrides."""
    result = dict(base)
    for key, value in override.items():
        if isinstance(value, dict) and isinstance(result.get(key), dict):
            result[key] = merge_mappings(dict(result[key]), value)
        else:
            result[key] = value
    return result


def merge_prompt_config(base: dict[str, Any], override: dict[str, Any]) -> dict[str, Any]:
    """Merge a small, output-run-specific prompt overlay safely.

    Nested objects are merged, while prompt lists are appended in stable
    de-duplicated order.  This lets a finite remediation run add a narrowly
    justified prompt without copying or silently editing the frozen base
    configuration used by a control run.  Scalar policy values remain an
    explicit override, so their provenance can be locked by the Slurm wrapper.
    """
    result = dict(base)
    for key, value in override.items():
        current = result.get(key)
        if isinstance(current, dict) and isinstance(value, dict):
            result[key] = merge_prompt_config(dict(current), value)
        elif isinstance(current, list) and isinstance(value, list):
            result[key] = list(dict.fromkeys([*current, *value]))
        else:
            result[key] = value
    return result


def float_or_default(value: Any, default: float) -> float:
    try:
        return float(value)
    except (TypeError, ValueError):
        return default


def normalized_float_mapping(value: Any) -> dict[str, float]:
    mapping = as_mapping(value)
    return {
        str(key).strip().lower(): float_or_default(item, 0.0)
        for key, item in mapping.items()
        if str(key).strip()
    }


def normalized_override_mapping(value: Any) -> dict[str, dict[str, Any]]:
    mapping = as_mapping(value)
    return {
        str(key).strip().lower(): as_mapping(item)
        for key, item in mapping.items()
        if str(key).strip() and isinstance(item, dict)
    }


def resolve_obstacle_evidence_policy(
    config: dict[str, Any], backend_name: str
) -> dict[str, Any]:
    """Compile a serializable trusted-obstacle-evidence policy.

    New configs may use ``obstacle_evidence_policy`` and optional
    ``obstacle_evidence_policy_by_backend.<backend>``.  The aliases
    ``obstacle_all_detected_policy`` / ``..._by_backend`` are accepted for
    early experiment configs.  In a legacy config, the existing
    ``obstacle_maximum_area_ratio`` map becomes the per-prompt area cap,
    confidence filtering remains disabled, and a conservative route-like
    component guard protects downstream hidden completion from broad scene
    surfaces.
    """
    base = as_mapping(
        config.get("obstacle_evidence_policy", config.get("obstacle_all_detected_policy", {}))
    )
    by_backend = as_mapping(
        config.get(
            "obstacle_evidence_policy_by_backend",
            config.get("obstacle_all_detected_policy_by_backend", {}),
        )
    )
    backend_override = as_mapping(by_backend.get(backend_name))
    raw = merge_mappings(base, backend_override)
    legacy_maximums = normalized_float_mapping(config.get("obstacle_maximum_area_ratio"))
    maximums = dict(legacy_maximums)
    maximums.update(
        normalized_float_mapping(
            raw.get("maximum_area_ratio", raw.get("maximum_area_ratio_by_prompt", {}))
        )
    )
    prompt_overrides = normalized_override_mapping(
        raw.get("prompt_overrides", raw.get("prompts", {}))
    )
    confidence_by_prompt = normalized_float_mapping(
        raw.get("minimum_confidence_by_prompt", raw.get("minimum_prompt_confidence_by_prompt", {}))
    )
    route_like_raw = as_mapping(
        raw.get("route_like", raw.get("route_like_component_policy", {}))
    )
    # The default is intentionally conservative: broad horizontal/ground-like
    # components are rejected from trusted evidence, but remain visible in
    # obstacle_unfiltered_evidence.png for human recovery.
    route_like = {
        "enabled": bool(route_like_raw.get("enabled", True)),
        "minimum_bbox_width_ratio": float_or_default(
            route_like_raw.get("minimum_bbox_width_ratio"), 0.80
        ),
        "minimum_bbox_height_ratio": float_or_default(
            route_like_raw.get("minimum_bbox_height_ratio"), 0.12
        ),
        "minimum_component_area_ratio": float_or_default(
            route_like_raw.get("minimum_component_area_ratio"), 0.03
        ),
    }
    default_maximum = float_or_default(
        raw.get("default_maximum_area_ratio", maximums.get("default", 0.15)),
        0.15,
    )
    return {
        "version": OBSTACLE_EVIDENCE_POLICY_VERSION,
        "backend": backend_name,
        "source": {
            "global": "obstacle_evidence_policy" if base else "legacy_defaults",
            "backend_override_present": bool(backend_override),
            "legacy_obstacle_maximum_area_ratio_used": bool(legacy_maximums),
        },
        "minimum_component_area_pixels": max(
            1, int(float_or_default(raw.get("minimum_component_area_pixels"), 16.0))
        ),
        "minimum_component_area_ratio": max(
            0.0, float_or_default(raw.get("minimum_component_area_ratio"), 0.00002)
        ),
        "minimum_prompt_confidence": min(
            1.0,
            max(
                0.0,
                float_or_default(
                    raw.get(
                        "minimum_prompt_confidence",
                        raw.get("minimum_confidence", 0.0),
                    ),
                    0.0,
                ),
            ),
        ),
        "minimum_confidence_by_prompt": confidence_by_prompt,
        "default_maximum_area_ratio": min(1.0, max(0.0, default_maximum)),
        "maximum_area_ratio_by_prompt": maximums,
        "full_frame_area_ratio": min(
            1.0, max(0.0, float_or_default(raw.get("full_frame_area_ratio"), 0.90))
        ),
        "prompt_overrides": prompt_overrides,
        "route_like": route_like,
    }


def evidence_minimum_component_area(policy: dict[str, Any], image_area: int) -> int:
    return max(
        int(policy["minimum_component_area_pixels"]),
        int(round(image_area * float(policy["minimum_component_area_ratio"]))),
    )


def evidence_prompt_limits(prompt: str, policy: dict[str, Any]) -> tuple[float, float]:
    key = prompt.strip().lower()
    override = as_mapping(policy["prompt_overrides"].get(key))
    confidence = float_or_default(
        override.get(
            "minimum_confidence",
            policy["minimum_confidence_by_prompt"].get(key, policy["minimum_prompt_confidence"]),
        ),
        float(policy["minimum_prompt_confidence"]),
    )
    maximum_area = float_or_default(
        override.get(
            "maximum_area_ratio",
            policy["maximum_area_ratio_by_prompt"].get(
                key, policy["default_maximum_area_ratio"]
            ),
        ),
        float(policy["default_maximum_area_ratio"]),
    )
    return min(1.0, max(0.0, confidence)), min(1.0, max(0.0, maximum_area))


def save_mask(path: Path, mask: np.ndarray) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    Image.fromarray(mask.astype(np.uint8) * 255, mode="L").save(path)


def save_red_mask(path: Path, mask: np.ndarray) -> None:
    """Save an RGB review visualization while preserving obstacle.png as binary."""
    path.parent.mkdir(parents=True, exist_ok=True)
    color = np.zeros((*mask.shape, 3), dtype=np.uint8)
    color[mask] = (230, 45, 45)
    Image.fromarray(color, mode="RGB").save(path)


def remove_small_components(mask: np.ndarray, minimum_area: int) -> np.ndarray:
    count, labels, stats, _ = cv2.connectedComponentsWithStats(mask.astype(np.uint8), 8)
    output = np.zeros_like(mask, dtype=bool)
    for index in range(1, count):
        if int(stats[index, cv2.CC_STAT_AREA]) >= minimum_area:
            output |= labels == index
    return output


def retain_core_connected_union(union: np.ndarray, core: np.ndarray, minimum_area: int) -> np.ndarray:
    union = remove_small_components(union, minimum_area)
    count, labels, stats, _ = cv2.connectedComponentsWithStats(union.astype(np.uint8), 8)
    output = np.zeros_like(union, dtype=bool)
    for index in range(1, count):
        component = labels == index
        if int(stats[index, cv2.CC_STAT_AREA]) >= minimum_area and np.any(component & core):
            output |= component
    return output


def retain_nearby_obstacles(
    obstacle: np.ndarray,
    target: np.ndarray,
    minimum_area: int,
    maximum_area_ratio: float = 0.15,
) -> np.ndarray:
    """Keep compact obstacle components spatially adjacent to the target surface."""
    distance = max(9, int(round(min(target.shape) * 0.035)) | 1)
    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (distance, distance))
    nearby = cv2.dilate(target.astype(np.uint8), kernel).astype(bool)
    count, labels, stats, _ = cv2.connectedComponentsWithStats(obstacle.astype(np.uint8), 8)
    output = np.zeros_like(obstacle, dtype=bool)
    image_area = obstacle.size
    for index in range(1, count):
        area = int(stats[index, cv2.CC_STAT_AREA])
        component = labels == index
        if (
            area >= minimum_area
            and area / image_area <= maximum_area_ratio
            and np.any(component & nearby)
        ):
            output |= component
    return output


def retain_compact_obstacles(
    obstacle: np.ndarray,
    minimum_area: int,
    maximum_area_ratio: float = 0.12,
) -> np.ndarray:
    """Keep compact scene obstacles without requiring target adjacency."""
    count, labels, stats, _ = cv2.connectedComponentsWithStats(obstacle.astype(np.uint8), 8)
    output = np.zeros_like(obstacle, dtype=bool)
    image_area = obstacle.size
    for index in range(1, count):
        area = int(stats[index, cv2.CC_STAT_AREA])
        if minimum_area <= area <= image_area * maximum_area_ratio:
            output |= labels == index
    return output


def route_like_component_count(mask: np.ndarray, route_policy: dict[str, Any]) -> int:
    """Count broad ground/route-like components that must not seed hidden masks."""
    if not bool(route_policy.get("enabled", False)) or not np.any(mask):
        return 0
    count, _, stats, _ = cv2.connectedComponentsWithStats(mask.astype(np.uint8), 8)
    height, width = mask.shape
    route_like = 0
    for index in range(1, count):
        area = int(stats[index, cv2.CC_STAT_AREA])
        width_ratio = float(stats[index, cv2.CC_STAT_WIDTH] / max(width, 1))
        height_ratio = float(stats[index, cv2.CC_STAT_HEIGHT] / max(height, 1))
        area_ratio = float(area / max(mask.size, 1))
        if (
            width_ratio >= float(route_policy["minimum_bbox_width_ratio"])
            and height_ratio >= float(route_policy["minimum_bbox_height_ratio"])
            and area_ratio >= float(route_policy["minimum_component_area_ratio"])
        ):
            route_like += 1
    return route_like


def trusted_obstacle_prompt_evidence(
    *,
    prompt: str,
    group: str,
    prediction: PromptPrediction,
    minimum_component_area: int,
    policy: dict[str, Any],
) -> tuple[np.ndarray, np.ndarray, dict[str, Any]]:
    """Return all cleaned evidence and the subset safe to seed hidden completion.

    The first mask is intentionally high recall: it only removes microscopic
    components.  The second is accepted only if its prompt confidence and
    image-area sanity checks pass and it is not route-like.  Rejecting at the
    prompt level makes the audit rationale easy to understand and guarantees
    that a broad false positive cannot leak into ``obstacle_all_detected``.
    """
    cleaned = remove_small_components(prediction.mask, minimum_component_area)
    minimum_confidence, maximum_area_ratio = evidence_prompt_limits(prompt, policy)
    area_ratio = float(cleaned.mean())
    reasons: list[str] = []
    if not np.any(cleaned):
        reasons.append("empty_after_small_component_removal")
    if float(prediction.confidence) < minimum_confidence:
        reasons.append("confidence_below_minimum")
    if area_ratio > maximum_area_ratio:
        reasons.append("area_ratio_exceeds_prompt_maximum")
    if area_ratio > float(policy["full_frame_area_ratio"]):
        reasons.append("area_ratio_exceeds_full_frame_guard")
    route_like_count = route_like_component_count(cleaned, policy["route_like"])
    if route_like_count:
        reasons.append("route_like_component_detected")
    accepted = not reasons
    trusted = cleaned if accepted else np.zeros_like(cleaned)
    decision = {
        "prompt": prompt,
        "group": group,
        "confidence": round(float(prediction.confidence), 6),
        "instance_count": int(prediction.instance_count),
        "minimum_confidence": round(minimum_confidence, 6),
        "maximum_area_ratio": round(maximum_area_ratio, 6),
        "cleaned_pixels": int(cleaned.sum()),
        "cleaned_area_ratio": round(area_ratio, 6),
        "route_like_component_count": route_like_count,
        "decision": "accepted" if accepted else "rejected",
        "reasons": reasons,
        "trusted_pixels": int(trusted.sum()),
    }
    return cleaned, trusted, decision


def area_plausibility(area_ratio: float, minimum: float, maximum: float) -> float:
    if area_ratio <= 0 or area_ratio < minimum / 2 or area_ratio > min(1.0, maximum * 1.2):
        return 0.0
    if minimum <= area_ratio <= maximum:
        return 1.0
    if area_ratio < minimum:
        return max(0.0, area_ratio / minimum)
    return max(0.0, 1.0 - (area_ratio - maximum) / max(1.0 - maximum, 1e-6))


@dataclass
class PromptPrediction:
    mask: np.ndarray
    confidence: float
    instance_count: int
    boxes: list[list[float]]


class Sam3Backend:
    name = "sam3"

    def __init__(self, args: argparse.Namespace):
        repo = Path(args.sam3_repo).resolve()
        checkpoint = Path(args.sam3_checkpoint).resolve()
        if not repo.is_dir():
            raise FileNotFoundError(f"SAM 3 repository not found: {repo}")
        if not checkpoint.is_file():
            raise FileNotFoundError(
                f"SAM 3 checkpoint not found: {checkpoint}. Accept the model terms and download sam3.pt first."
            )
        sys.path.insert(0, str(repo))
        import torch
        from sam3.model.sam3_image_processor import Sam3Processor
        from sam3.model_builder import build_sam3_image_model

        self.torch = torch
        self.device = args.device
        self.autocast_device = str(args.device).split(":", 1)[0]
        self.autocast_dtype = torch.bfloat16
        model = build_sam3_image_model(
            device=args.device,
            checkpoint_path=str(checkpoint),
            load_from_HF=False,
            compile=args.compile,
        )
        self.processor = Sam3Processor(
            model,
            device=args.device,
            confidence_threshold=args.confidence_threshold,
        )
        self.state: dict[str, Any] | None = None

    def begin_image(self, image: Image.Image) -> None:
        with self.torch.autocast(
            device_type=self.autocast_device,
            dtype=self.autocast_dtype,
            enabled=self.autocast_device == "cuda",
        ):
            self.state = self.processor.set_image(image)

    def predict(self, prompt: str) -> PromptPrediction:
        if self.state is None:
            raise RuntimeError("begin_image must be called before predict")
        with self.torch.autocast(
            device_type=self.autocast_device,
            dtype=self.autocast_dtype,
            enabled=self.autocast_device == "cuda",
        ):
            output = self.processor.set_text_prompt(prompt, self.state)
        masks_tensor = output["masks"]
        scores_tensor = output["scores"]
        boxes_tensor = output["boxes"]
        if masks_tensor.numel() == 0:
            shape = (int(output["original_height"]), int(output["original_width"]))
            return PromptPrediction(np.zeros(shape, dtype=bool), 0.0, 0, [])
        masks = masks_tensor.detach().cpu().numpy().astype(bool)
        while masks.ndim > 3 and masks.shape[1] == 1:
            masks = masks[:, 0]
        union = np.any(masks, axis=0)
        scores = scores_tensor.detach().float().cpu().numpy()
        boxes = boxes_tensor.detach().float().cpu().numpy().tolist()
        return PromptPrediction(union, float(scores.max()), int(len(scores)), boxes)


class GroundedSamBackend:
    name = "grounded_sam"

    def __init__(self, args: argparse.Namespace):
        repo = Path(args.grounded_sam_repo).resolve()
        config = Path(args.grounding_dino_config).resolve()
        dino_checkpoint = Path(args.grounding_dino_checkpoint).resolve()
        sam_checkpoint = Path(args.sam_checkpoint).resolve()
        for path in (repo, config, dino_checkpoint, sam_checkpoint):
            if not path.exists():
                raise FileNotFoundError(f"Required Grounded-SAM path not found: {path}")
        sys.path[:0] = [str(repo), str(repo / "GroundingDINO"), str(repo / "segment_anything")]
        import torch
        import GroundingDINO.groundingdino.datasets.transforms as T
        from GroundingDINO.groundingdino.models import build_model
        from GroundingDINO.groundingdino.util.slconfig import SLConfig
        from GroundingDINO.groundingdino.util.utils import clean_state_dict
        from segment_anything import SamPredictor, sam_model_registry

        model_args = SLConfig.fromfile(str(config))
        model_args.device = args.device
        if args.bert_path:
            model_args.bert_base_uncased_path = str(Path(args.bert_path).resolve())
        self.torch = torch
        self.transforms = T.Compose(
            [
                T.RandomResize([800], max_size=1333),
                T.ToTensor(),
                T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
            ]
        )
        self.device = args.device
        self.box_threshold = args.box_threshold
        self.model = build_model(model_args)
        checkpoint = torch.load(str(dino_checkpoint), map_location="cpu", weights_only=False)
        self.model.load_state_dict(clean_state_dict(checkpoint["model"]), strict=False)
        self.model.to(args.device).eval()
        sam = sam_model_registry[args.sam_model_type](checkpoint=str(sam_checkpoint)).to(args.device)
        self.predictor = SamPredictor(sam)
        self.image_pil: Image.Image | None = None
        self.image_tensor = None
        self.image_rgb: np.ndarray | None = None

    def begin_image(self, image: Image.Image) -> None:
        self.image_pil = image
        self.image_rgb = np.asarray(image)
        self.image_tensor, _ = self.transforms(image, None)
        self.predictor.set_image(self.image_rgb)

    def predict(self, prompt: str) -> PromptPrediction:
        if self.image_pil is None or self.image_tensor is None or self.image_rgb is None:
            raise RuntimeError("begin_image must be called before predict")
        caption = prompt.lower().strip()
        if not caption.endswith("."):
            caption += "."
        with self.torch.inference_mode():
            outputs = self.model(self.image_tensor[None].to(self.device), captions=[caption])
        logits = outputs["pred_logits"].detach().cpu().sigmoid()[0]
        boxes = outputs["pred_boxes"].detach().cpu()[0]
        confidence = logits.max(dim=1).values
        keep = confidence > self.box_threshold
        boxes = boxes[keep]
        confidence = confidence[keep]
        height, width = self.image_rgb.shape[:2]
        if boxes.shape[0] == 0:
            return PromptPrediction(np.zeros((height, width), dtype=bool), 0.0, 0, [])
        scale = self.torch.tensor([width, height, width, height], dtype=boxes.dtype)
        boxes = boxes * scale
        boxes[:, :2] -= boxes[:, 2:] / 2
        boxes[:, 2:] += boxes[:, :2]
        transformed = self.predictor.transform.apply_boxes_torch(boxes, (height, width)).to(self.device)
        with self.torch.inference_mode():
            masks, sam_scores, _ = self.predictor.predict_torch(
                point_coords=None,
                point_labels=None,
                boxes=transformed,
                multimask_output=False,
            )
        masks_np = masks[:, 0].detach().cpu().numpy().astype(bool)
        combined_confidence = float(
            math.sqrt(float(confidence.max()) * float(sam_scores.detach().cpu().max()))
        )
        return PromptPrediction(
            np.any(masks_np, axis=0),
            combined_confidence,
            int(len(masks_np)),
            boxes.tolist(),
        )


def build_backend(args: argparse.Namespace):
    if args.backend == "sam3":
        return Sam3Backend(args)
    if args.backend == "grounded_sam":
        return GroundedSamBackend(args)
    raise ValueError(f"Unsupported backend: {args.backend}")


def prompt_consensus(
    predictions: list[PromptPrediction], minimum_area: int
) -> tuple[np.ndarray, np.ndarray, np.ndarray, dict[str, Any]]:
    if not predictions:
        raise ValueError("At least one prompt prediction is required")
    masks = np.stack([prediction.mask for prediction in predictions])
    votes = masks.sum(axis=0)
    union = votes > 0
    majority = max(2, math.ceil(len(predictions) / 2)) if len(predictions) > 1 else 1
    core = votes >= majority
    fallback = False
    if not np.any(core) and np.any(union):
        fallback = True
        best_index = int(np.argmax([prediction.confidence for prediction in predictions]))
        core = predictions[best_index].mask.copy()
    candidate = retain_core_connected_union(union, core, minimum_area)
    core = remove_small_components(core & candidate, minimum_area)
    uncertain = candidate & ~core
    union_area = int(candidate.sum())
    agreement = float(core.sum() / union_area) if union_area else 0.0
    confidences = [prediction.confidence for prediction in predictions]
    metadata = {
        "majority_vote": majority,
        "fallback_to_best_prompt": fallback,
        "agreement": round(agreement, 6),
        "mean_prompt_confidence": round(float(np.mean(confidences)), 6),
        "maximum_prompt_confidence": round(float(np.max(confidences)), 6),
        "detected_prompt_count": sum(value > 0 for value in confidences),
    }
    return candidate, core, uncertain, metadata


def overlay_masks(
    image: np.ndarray, target: np.ndarray, obstacle: np.ndarray, uncertain: np.ndarray
) -> np.ndarray:
    overlay = image.astype(np.float32).copy()
    colors = [
        (target, np.asarray([30, 210, 70], dtype=np.float32), 0.38),
        (uncertain, np.asarray([250, 210, 30], dtype=np.float32), 0.55),
        # Draw obstacles last so target/uncertainty colors never hide red occluders.
        (obstacle, np.asarray([230, 45, 45], dtype=np.float32), 0.58),
    ]
    for mask, color, alpha in colors:
        overlay[mask] = overlay[mask] * (1.0 - alpha) + color * alpha
    return np.clip(overlay, 0, 255).astype(np.uint8)


def score_category(
    mask: np.ndarray,
    consensus: dict[str, Any],
    thresholds: dict[str, float],
    category: str,
) -> tuple[float, dict[str, float]]:
    area_ratio = float(mask.mean())
    plausibility = area_plausibility(
        area_ratio,
        thresholds["minimum_target_area_ratio"],
        thresholds["maximum_target_area_ratio"],
    )
    bottom_contact = float(mask[int(mask.shape[0] * 0.8) :, :].mean() > 0.005)
    confidence = float(consensus["mean_prompt_confidence"])
    agreement = float(consensus["agreement"])
    evidence = 0.55 * confidence + 0.35 * agreement + 0.10 * bottom_contact
    score = evidence * (0.20 + 0.80 * plausibility)
    if category == "walkway":
        score *= 0.94  # prevent the generic class from dominating specific structures
    signals = {
        "area_ratio": round(area_ratio, 6),
        "area_plausibility": round(plausibility, 6),
        "bottom_contact": bottom_contact,
        "score": round(score, 6),
    }
    return score, signals


def predict_prompt_set(
    backend,
    prompts: list[str],
    output_dir: Path,
    minimum_component_area: int,
    save_prompt_masks: bool = True,
) -> tuple[np.ndarray, np.ndarray, np.ndarray, dict[str, Any]]:
    predictions = []
    prompt_metadata = []
    for prompt in prompts:
        prediction = backend.predict(prompt)
        predictions.append(prediction)
        if save_prompt_masks:
            save_mask(output_dir / f"{safe_name(prompt)}.png", prediction.mask)
        prompt_metadata.append(
            {
                "prompt": prompt,
                "confidence": round(prediction.confidence, 6),
                "instance_count": prediction.instance_count,
                "boxes_xyxy": prediction.boxes,
            }
        )
    candidate, core, uncertain, consensus = prompt_consensus(
        predictions, minimum_component_area
    )
    consensus["prompts"] = prompt_metadata
    return candidate, core, uncertain, consensus


def process_sample(
    row: dict[str, Any],
    backend,
    config: dict[str, Any],
    output_root: Path,
    args: argparse.Namespace,
) -> dict[str, Any]:
    sample_id = row["sample_id"]
    sample_dir = output_root / "samples" / sample_id
    quality_path = sample_dir / "quality.json"
    if args.resume:
        quality = load_resume_quality(sample_dir, sample_id, backend.name)
        if quality is not None:
            return quality

    # Phone photos commonly store the camera pixels in landscape orientation
    # and rely on EXIF to display them upright.  Normalize that orientation
    # before segmentation so every downstream mask uses the displayed frame.
    image = ImageOps.exif_transpose(Image.open(row["image_path"])).convert("RGB")
    rgb = np.asarray(image)
    image_area = rgb.shape[0] * rgb.shape[1]
    minimum_component_area = max(args.minimum_component_area, int(image_area * 0.0002))
    backend.begin_image(image)

    requested_categories = args.categories or list(config["categories"])
    fixed_category = row.get("category")
    if fixed_category:
        requested_categories = [fixed_category]
        if args.verify_manifest_category and not row.get("category_reviewed", False):
            requested_categories.extend(
                config.get("category_confusions", {}).get(fixed_category, [])
            )
            requested_categories = list(dict.fromkeys(requested_categories))

    category_results: dict[str, dict[str, Any]] = {}
    masks: dict[str, tuple[np.ndarray, np.ndarray, np.ndarray]] = {}
    thresholds = config["quality_thresholds"]
    save_full_artifacts = args.artifact_level == "full"
    for category in requested_categories:
        if category not in config["categories"]:
            raise ValueError(f"Unknown category in pilot manifest: {category}")
        prompt_dir = sample_dir / "prompts" / category
        candidate, core, uncertain, consensus = predict_prompt_set(
            backend,
            config["categories"][category]["prompts"],
            prompt_dir,
            minimum_component_area,
            save_prompt_masks=save_full_artifacts,
        )
        # --- per-category area cap (prevents whole-image over-segmentation) ---
        cat_max_area = config["categories"][category].get(
            "max_area_ratio",
            config.get("tactile_paving_max_area_ratio", None),
        )
        candidate_area_ratio = float(candidate.mean()) if candidate.any() else 0.0
        if cat_max_area is not None and candidate_area_ratio > float(cat_max_area):
            consensus["area_cap_rejected"] = True
            consensus["area_cap_ratio"] = round(candidate_area_ratio, 6)
            consensus["area_cap_limit"] = float(cat_max_area)
            candidate = np.zeros_like(candidate)
            core = np.zeros_like(core)
            uncertain = np.zeros_like(uncertain)
        else:
            consensus["area_cap_rejected"] = False
        # --- end area cap ---
        score, signals = score_category(candidate, consensus, thresholds, category)
        if save_full_artifacts:
            save_mask(sample_dir / "categories" / f"{category}_candidate.png", candidate)
            save_mask(sample_dir / "categories" / f"{category}_core.png", core)
            save_mask(sample_dir / "categories" / f"{category}_uncertain.png", uncertain)
        masks[category] = (candidate, core, uncertain)
        category_results[category] = {"consensus": consensus, **signals}

    best_scored_category = max(
        category_results, key=lambda key: category_results[key]["score"]
    )
    # A mask score is useful for flagging semantic disagreement, but it is not
    # reliable enough to silently overwrite a manifest category. Human review
    # overrides remain authoritative.
    selected_category = fixed_category or best_scored_category
    category_disagreement = bool(
        fixed_category and best_scored_category != fixed_category
    )
    target, target_core, target_uncertain = masks[selected_category]

    prompt_resolution = obstacle_prompt_resolution(
        config, selected_category, backend.name
    )
    obstacle_predictions: list[tuple[str, PromptPrediction]] = []
    obstacle_metadata = []
    for prompt in prompt_resolution["prompts"]:
        prediction = backend.predict(prompt)
        obstacle_predictions.append((prompt, prediction))
        if save_full_artifacts:
            save_mask(
                sample_dir / "prompts" / "obstacles" / f"{safe_name(prompt)}.png",
                prediction.mask,
            )
        obstacle_metadata.append(
            {
                "prompt": prompt,
                "group": "dynamic_or_compact",
                "confidence": round(prediction.confidence, 6),
                "instance_count": prediction.instance_count,
                "boxes_xyxy": prediction.boxes,
            }
        )
    maximum_ratios = config.get("obstacle_maximum_area_ratio", {})
    filtered_dynamic_obstacles = []
    for prompt, prediction in obstacle_predictions:
        cleaned = remove_small_components(prediction.mask, minimum_component_area)
        maximum_area_ratio = float(
            maximum_ratios.get(prompt, maximum_ratios.get("default", 0.15))
        )
        filtered_dynamic_obstacles.append(
            retain_nearby_obstacles(
                cleaned,
                target,
                minimum_component_area,
                maximum_area_ratio=maximum_area_ratio,
            )
        )
    dynamic_obstacle = (
        np.any(np.stack(filtered_dynamic_obstacles), axis=0)
        if filtered_dynamic_obstacles
        else np.zeros_like(target)
    )

    barrier_predictions: list[tuple[str, PromptPrediction]] = []
    barrier_backends = config.get("barrier_obstacle_backends", ["sam3"])
    barrier_prompts = (
        config.get("barrier_obstacle_prompts", [])
        if backend.name in barrier_backends
        else []
    )
    for prompt in barrier_prompts:
        prediction = backend.predict(prompt)
        barrier_predictions.append((prompt, prediction))
        if save_full_artifacts:
            save_mask(
                sample_dir / "prompts" / "obstacles" / f"{safe_name(prompt)}.png",
                prediction.mask,
            )
        obstacle_metadata.append(
            {
                "prompt": prompt,
                "group": "barrier_consensus",
                "confidence": round(prediction.confidence, 6),
                "instance_count": prediction.instance_count,
                "boxes_xyxy": prediction.boxes,
            }
        )
    if barrier_predictions:
        barrier_votes = np.stack(
            [prediction.mask for _, prediction in barrier_predictions]
        ).sum(axis=0)
        required_votes = max(2, math.ceil(len(barrier_predictions) / 2))
        barrier_obstacle = barrier_votes >= required_votes
        barrier_obstacle = retain_compact_obstacles(
            barrier_obstacle, minimum_component_area
        )
    else:
        barrier_obstacle = np.zeros_like(target)

    # Keep two explicit evidence layers.  The unfiltered union is an audit
    # artifact: every prompt after only microscopic-component removal.  The
    # trusted union applies per-prompt confidence, prompt-area, full-frame,
    # and broad-route sanity checks before it can seed a later hidden proposal.
    # ``obstacle.png`` below intentionally keeps its original nearby policy.
    evidence_policy = resolve_obstacle_evidence_policy(config, backend.name)
    evidence_minimum_area = evidence_minimum_component_area(evidence_policy, image_area)
    unfiltered_evidence_masks: list[np.ndarray] = []
    trusted_evidence_masks: list[np.ndarray] = []
    obstacle_evidence_decisions: list[dict[str, Any]] = []
    for group, predictions in (
        ("dynamic_or_compact", obstacle_predictions),
        ("barrier_consensus", barrier_predictions),
    ):
        for prompt, prediction in predictions:
            unfiltered, trusted, decision = trusted_obstacle_prompt_evidence(
                prompt=prompt,
                group=group,
                prediction=prediction,
                minimum_component_area=evidence_minimum_area,
                policy=evidence_policy,
            )
            unfiltered_evidence_masks.append(unfiltered)
            trusted_evidence_masks.append(trusted)
            obstacle_evidence_decisions.append(decision)
    obstacle_unfiltered_evidence = (
        np.any(np.stack(unfiltered_evidence_masks), axis=0)
        if unfiltered_evidence_masks
        else np.zeros_like(target)
    )
    obstacle_all_detected = (
        np.any(np.stack(trusted_evidence_masks), axis=0)
        if trusted_evidence_masks
        else np.zeros_like(target)
    )
    obstacle = dynamic_obstacle | barrier_obstacle
    target = target & ~obstacle
    target_core = target_core & target
    target_uncertain = target_uncertain & target

    kernel_size = max(5, int(round(min(rgb.shape[:2]) * 0.012)) | 1)
    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size, kernel_size))
    target_dilated = cv2.dilate(target.astype(np.uint8), kernel).astype(bool)
    obstacle_dilated = cv2.dilate(obstacle.astype(np.uint8), kernel).astype(bool)
    adjacency = target_dilated & obstacle_dilated
    adjacency_ratio = float(adjacency.sum() / max(target.sum(), 1))

    save_mask(sample_dir / "target_visible_candidate.png", target)
    save_mask(sample_dir / "target_visible_core.png", target_core)
    save_mask(sample_dir / "target_uncertain.png", target_uncertain)
    save_mask(sample_dir / "obstacle_unfiltered_evidence.png", obstacle_unfiltered_evidence)
    save_mask(sample_dir / "obstacle_all_detected.png", obstacle_all_detected)
    save_mask(sample_dir / "obstacle.png", obstacle)
    save_red_mask(sample_dir / "obstacle_red.png", obstacle)
    Image.fromarray(overlay_masks(rgb, target_core, obstacle, target_uncertain)).save(
        sample_dir / "overlay.png", quality=92
    )

    selected = category_results[selected_category]
    score = float(selected["score"])
    target_area_ratio = float(target.mean())
    uncertainty_ratio = float(target_uncertain.sum() / max(target.sum(), 1))
    obstacle_area_ratio = float(obstacle.mean())
    area_valid = (
        thresholds["minimum_target_area_ratio"]
        <= target_area_ratio
        <= thresholds["maximum_target_area_ratio"]
    )
    if not area_valid:
        review_status = "reject_or_reprompt"
    elif (
        score >= thresholds["auto_accept_score"]
        and not selected["consensus"]["fallback_to_best_prompt"]
        and uncertainty_ratio <= 0.35
        and obstacle_area_ratio <= 0.20
    ):
        review_status = "candidate_accept_after_visual_review"
    elif score >= thresholds["manual_review_score"]:
        review_status = "manual_review"
    else:
        review_status = "reject_or_reprompt"
    if (
        category_disagreement
        and not row.get("category_reviewed", False)
        and review_status == "candidate_accept_after_visual_review"
    ):
        review_status = "manual_review"
    quality = {
        "sample_id": sample_id,
        "image_path": row["image_path"],
        "backend": backend.name,
        "selected_category": selected_category,
        "category_verification": {
            "manifest_category": fixed_category,
            "best_scored_category": best_scored_category,
            "category_disagreement": category_disagreement,
            "category_reviewed": bool(row.get("category_reviewed", False)),
        },
        "quality_score": round(score, 6),
        "review_status": review_status,
        "target_area_ratio": round(target_area_ratio, 6),
        "target_core_ratio": round(float(target_core.mean()), 6),
        "uncertainty_ratio_within_target": round(
            uncertainty_ratio, 6
        ),
        "obstacle_area_ratio": round(obstacle_area_ratio, 6),
        "obstacle_unfiltered_evidence_pixels": int(obstacle_unfiltered_evidence.sum()),
        "obstacle_unfiltered_evidence_area_ratio": round(
            float(obstacle_unfiltered_evidence.mean()), 6
        ),
        "obstacle_all_detected_pixels": int(obstacle_all_detected.sum()),
        "obstacle_all_detected_area_ratio": round(float(obstacle_all_detected.mean()), 6),
        "obstacle_all_detected_policy": "sanitized_high_recall_prompt_union_v2",
        "obstacle_evidence_policy_version": OBSTACLE_EVIDENCE_POLICY_VERSION,
        "obstacle_evidence_minimum_component_area_pixels": evidence_minimum_area,
        "obstacle_evidence_artifact_semantics": {
            "obstacle_unfiltered_evidence.png": (
                "union of every obstacle/barrier prompt after only small-component removal; "
                "audit evidence only, never direct hidden-completion support"
            ),
            "obstacle_all_detected.png": (
                "high-recall prompt union after confidence, per-prompt area, full-frame, "
                "and route-like sanity filters; candidate evidence for later hidden completion"
            ),
            "obstacle.png": (
                "unchanged nearby/compact obstacle proposal used for visible-mask cleanup; "
                "not the high-recall evidence layer"
            ),
        },
        "obstacle_prompt_resolution": prompt_resolution,
        "obstacle_evidence_policy": evidence_policy,
        "obstacle_evidence_prompt_decisions": obstacle_evidence_decisions,
        "obstacle_evidence_decision_counts": dict(
            sorted(Counter(item["decision"] for item in obstacle_evidence_decisions).items())
        ),
        "barrier_obstacle_area_ratio": round(float(barrier_obstacle.mean()), 6),
        "target_obstacle_adjacency_ratio": round(adjacency_ratio, 6),
        "category_results": category_results,
        "obstacle_prompts": obstacle_metadata,
        "automatic_mask_is_ground_truth": False,
        "artifact_level": args.artifact_level,
    }
    json_dump(quality_path, quality)
    return quality


def build_html_report(output_dir: Path, results: list[dict[str, Any]]) -> None:
    rows = []
    for result in sorted(results, key=lambda item: item.get("quality_score", -1), reverse=True):
        sample_id = result["sample_id"]
        overlay = f"samples/{sample_id}/overlay.png"
        error = result.get("error")
        details = html.escape(error) if error else (
            f'{html.escape(result["selected_category"])} | '
            f'score={result["quality_score"]:.3f} | '
            f'{html.escape(result["review_status"])}'
        )
        rows.append(
            f'<article><img src="{html.escape(overlay)}" loading="lazy">'
            f'<div><strong>{html.escape(sample_id)}</strong><br>{details}</div></article>'
        )
    document = """<!doctype html><meta charset="utf-8"><title>Accessibility mask pilot</title>
<style>body{font-family:sans-serif;margin:20px}main{display:grid;grid-template-columns:repeat(3,minmax(0,1fr));gap:16px}article{border:1px solid #bbb;padding:8px}img{width:100%;height:auto}div{margin-top:6px;font-size:14px}</style>
<h1>Accessibility mask pilot</h1><p>Green: consensus core; yellow: uncertain target; red: obstacle. All masks require visual review.</p><main>"""
    document += "\n".join(rows) + "</main>"
    (output_dir / "report.html").write_text(document, encoding="utf-8")


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--pilot-manifest", default="output/accessibility_clip_screen/pilot_manifest.jsonl")
    parser.add_argument("--image", default=None, help="Single-image mode input image; bypasses --pilot-manifest.")
    parser.add_argument("--single-sample-id", default=None, help="Stable output ID used with --image.")
    parser.add_argument(
        "--single-category",
        choices=["curb_cut", "ramp", "stairs", "tactile_paving", "walkway"],
        default=None,
        help="Known category for --image. Omit to score all configured categories.",
    )
    parser.add_argument("--prompt-config", default="configs/accessibility_mask_prompts.json")
    parser.add_argument(
        "--prompt-config-override",
        default=None,
        help=(
            "Optional JSON overlay merged into --prompt-config for this output-only run. "
            "Nested prompt lists are appended and de-duplicated; the base config is never edited."
        ),
    )
    parser.add_argument("--output-dir", default="output/accessibility_mask_proposals/sam3")
    parser.add_argument("--backend", choices=["sam3", "grounded_sam"], default="sam3")
    parser.add_argument("--device", default="cuda")
    parser.add_argument("--categories", nargs="*", default=None)
    parser.add_argument("--limit", type=int, default=0)
    parser.add_argument("--minimum-component-area", type=int, default=128)
    parser.add_argument("--confidence-threshold", type=float, default=0.38)
    parser.add_argument("--compile", action="store_true")
    parser.add_argument("--resume", action="store_true")
    parser.add_argument(
        "--artifact-level",
        choices=["full", "standard"],
        default="full",
        help="standard omits regenerable per-prompt/category PNGs for large batch runs.",
    )
    parser.add_argument(
        "--verify-manifest-category",
        action="store_true",
        help="Compare configured confusable categories unless category_reviewed=true.",
    )
    parser.add_argument(
        "--allow-test-split",
        action="store_true",
        help=(
            "Explicitly allow frozen test RGB rows for a pre-registered front-end "
            "inference run. Default behavior remains validation-only."
        ),
    )
    parser.add_argument("--sam3-repo", default="repos/sam3")
    parser.add_argument("--sam3-checkpoint", default="weights/sam3/sam3.pt")
    parser.add_argument(
        "--grounded-sam-repo",
        default="repos/grounded-sam-compat",
    )
    parser.add_argument(
        "--grounding-dino-config",
        default="repos/grounded-sam-compat/GroundingDINO/groundingdino/config/GroundingDINO_SwinT_OGC.py",
    )
    parser.add_argument(
        "--grounding-dino-checkpoint",
        default="weights/grounded-sam/groundingdino_swint_ogc.pth",
    )
    parser.add_argument(
        "--sam-checkpoint",
        default="weights/grounded-sam/sam_vit_h_4b8939.pth",
    )
    parser.add_argument("--sam-model-type", choices=["vit_h", "vit_l", "vit_b"], default="vit_h")
    parser.add_argument("--bert-path", default=None)
    parser.add_argument("--box-threshold", type=float, default=0.28)
    return parser


def main() -> int:
    args = build_parser().parse_args()
    config_path = Path(args.prompt_config).resolve()
    config_override_path = (
        Path(args.prompt_config_override).resolve()
        if args.prompt_config_override
        else None
    )
    output_dir = Path(args.output_dir).resolve()
    output_dir.mkdir(parents=True, exist_ok=True)
    with config_path.open("r", encoding="utf-8") as handle:
        config = json.load(handle)
    if not isinstance(config, dict):
        raise ValueError(f"prompt config must be a JSON object: {config_path}")
    if config_override_path is not None:
        with config_override_path.open("r", encoding="utf-8") as handle:
            override = json.load(handle)
        if not isinstance(override, dict):
            raise ValueError(f"prompt config override must be a JSON object: {config_override_path}")
        config = merge_prompt_config(config, override)
    if args.image:
        image_path = Path(args.image).expanduser().resolve()
        if not image_path.is_file():
            raise FileNotFoundError(image_path)
        sample_id = args.single_sample_id or safe_name(image_path.stem) or "single_image"
        rows = [
            {
                "sample_id": sample_id,
                "image_path": str(image_path),
                "category": args.single_category,
                "category_reviewed": bool(args.single_category),
            }
        ]
    else:
        pilot_path = Path(args.pilot_manifest).resolve()
        rows = read_jsonl(pilot_path)
    if args.limit > 0:
        rows = rows[: args.limit]
    allowed_splits = {None, "validation"}
    if args.allow_test_split:
        allowed_splits.add("test")
    invalid_splits = sorted(
        {str(row.get("benchmark_split")) for row in rows if row.get("benchmark_split") not in allowed_splits}
    )
    if invalid_splits:
        raise ValueError(
            "Proposal manifest contains unsupported benchmark split(s): "
            f"{invalid_splits}. Pass --allow-test-split only for an explicitly "
            "registered test front-end inference run."
        )
    test_accessed = any(row.get("benchmark_split") == "test" for row in rows)
    backend = None
    initialization_error = None
    try:
        backend = build_backend(args)
        if backend.torch.cuda.is_available():
            backend.torch.cuda.reset_peak_memory_stats()
    except Exception as exc:
        traceback.print_exc()
        initialization_error = f"{type(exc).__name__}: {exc}"
    results = []
    if initialization_error is not None:
        for row in rows:
            results.append(
                {
                    "sample_id": row["sample_id"],
                    "image_path": row["image_path"],
                    "backend": args.backend,
                    "quality_score": -1.0,
                    "review_status": "error",
                    "error": f"model_initialization_failed: {initialization_error}",
                }
            )
        write_jsonl(output_dir / "results.jsonl", results)
    else:
        assert backend is not None
        for index, row in enumerate(rows, start=1):
            print(f"[{index}/{len(rows)}] {row['sample_id']}", flush=True)
            try:
                result = process_sample(row, backend, config, output_dir, args)
            except Exception as exc:
                traceback.print_exc()
                result = {
                    "sample_id": row["sample_id"],
                    "image_path": row["image_path"],
                    "backend": args.backend,
                    "quality_score": -1.0,
                    "review_status": "error",
                    "error": f"{type(exc).__name__}: {exc}",
                }
            results.append(result)
            write_jsonl(output_dir / "results.jsonl", results)
    build_html_report(output_dir, results)
    failed = sum(item.get("review_status") == "error" for item in results)
    peak_bytes = (
        int(backend.torch.cuda.max_memory_allocated())
        if backend is not None and backend.torch.cuda.is_available()
        else 0
    )
    summary = {
        "backend": args.backend,
        "sample_count": len(results),
        "completed": len(results) - failed,
        "failed": failed,
        "failure_rate": failed / len(results) if results else 1.0,
        "status_counts": dict(
            sorted(
                {
                    status: sum(item.get("review_status") == status for item in results)
                    for status in {item.get("review_status") for item in results}
                }.items()
            )
        ),
        "mean_quality_score": round(
            float(np.mean([item["quality_score"] for item in results if item["quality_score"] >= 0])),
            6,
        )
        if any(item["quality_score"] >= 0 for item in results)
        else None,
        "automatic_masks_are_ground_truth": False,
        "prompt_config": str(config_path),
        "prompt_config_override": str(config_override_path) if config_override_path else None,
        "model_initialization_error": initialization_error,
        "peak_gpu_memory_bytes": peak_bytes,
        "peak_gpu_memory_gib": peak_bytes / (1024**3),
        "environment": {
            "python": sys.version,
            "executable": sys.executable,
            "torch": getattr(getattr(backend, "torch", None), "__version__", None),
            "cuda_runtime": getattr(getattr(backend, "torch", None), "version", None).cuda
            if backend is not None
            else None,
            "gpu_name": backend.torch.cuda.get_device_name(0)
            if backend is not None and backend.torch.cuda.is_available()
            else None,
            "slurm_job_id": os.environ.get("SLURM_JOB_ID"),
        },
        "resolved_arguments": vars(args),
        "test_accessed": test_accessed,
        "test_access_explicitly_allowed": bool(args.allow_test_split),
    }
    json_dump(output_dir / "summary.json", summary)
    print(json.dumps(summary, ensure_ascii=False, indent=2))
    return 0 if all(item.get("review_status") != "error" for item in results) else 2


if __name__ == "__main__":
    raise SystemExit(main())