evaluation/tools/paired_bootstrap.py

"""Paired image-cluster bootstrap for offline Orienter result checks.

This tool is intentionally independent of pycocotools and semantic_matching.py:
semantic mode reads a frozen embedding cache directly and fails on cache miss,
so it never imports OpenAI/Zhipu clients or writes cache files.
"""

import argparse
import json
import math
import re
import statistics
from dataclasses import dataclass
from pathlib import Path
from typing import Callable, Dict, Iterable, List, Mapping, Optional, Sequence, Tuple

import numpy as np


SIMILARITY_THRESHOLD = 0.85


@dataclass(frozen=True)
class GroundTruth:
    image_ids: Tuple[int, ...]
    annotations_by_image: Mapping[int, Tuple[dict, ...]]
    category_names: Mapping[int, str]


@dataclass(frozen=True)
class ImageMetrics:
    precision: float
    recall: float
    f1: float
    support: int
    predictions: int


@dataclass(frozen=True)
class EvaluationResult:
    metrics: Mapping[str, float]
    per_image: Mapping[int, ImageMetrics]
    counts: Mapping[str, int]
    score_threshold: float


class SemanticMatcher:
    def __init__(
        self,
        cache_path: Path,
        *,
        similarity_threshold: float = SIMILARITY_THRESHOLD,
        similarity_mode: str = "raw_dot",
    ):
        self.cache_path = Path(cache_path)
        self.similarity_threshold = similarity_threshold
        self.similarity_mode = similarity_mode
        self.embeddings = self._load_cache(self.cache_path)
        self.dimension = self._embedding_dimension(self.embeddings.values())

    @staticmethod
    def standardize(category: object) -> str:
        text = re.sub(r"\d+$", "", str(category)).strip()
        text = text.replace("_", " ").lower().title()
        text = " ".join(text.split())
        if len(text.split()) == 2:
            text = "".join(text.split())
        return text

    @staticmethod
    def _embedding_dimension(values: Iterable[Sequence[float]]) -> Optional[int]:
        for value in values:
            if value is not None:
                return len(value)
        return None

    @classmethod
    def _load_cache(cls, path: Path) -> Dict[str, Tuple[float, ...]]:
        with path.open(encoding="utf-8") as file:
            payload = json.load(file)
        if not isinstance(payload, dict):
            raise ValueError(f"Embedding cache must be a JSON object: {path}")

        cache = {}
        expected_dim = None
        for key, value in payload.items():
            if value is None:
                continue
            if not isinstance(key, str):
                raise ValueError(f"Embedding cache key must be a string in {path}")
            if not isinstance(value, list) or not value:
                raise ValueError(f"Embedding cache value for {key!r} must be a non-empty list")
            if not all(isinstance(item, (int, float)) and math.isfinite(item) for item in value):
                raise ValueError(f"Embedding cache value for {key!r} must contain finite numbers")
            if expected_dim is None:
                expected_dim = len(value)
            elif len(value) != expected_dim:
                raise ValueError(
                    f"Embedding cache value for {key!r} has dimension {len(value)}, expected {expected_dim}"
                )
            cache[key] = tuple(float(item) for item in value)
        return cache

    def matches(self, gt_category: object, pred_category: object) -> bool:
        gt_label = self.standardize(gt_category)
        pred_label = self.standardize(pred_category)
        missing = [label for label in (gt_label, pred_label) if label not in self.embeddings]
        if missing:
            unique_missing = ", ".join(repr(label) for label in sorted(set(missing)))
            raise KeyError(
                "Missing frozen embedding cache entries: "
                f"{unique_missing}. Refusing to call an embedding API."
            )

        gt_vec = np.asarray(self.embeddings[gt_label], dtype=float)
        pred_vec = np.asarray(self.embeddings[pred_label], dtype=float)
        if self.similarity_mode == "raw_dot":
            similarity = float(np.dot(gt_vec, pred_vec))
        elif self.similarity_mode == "cosine":
            gt_norm = np.linalg.norm(gt_vec)
            pred_norm = np.linalg.norm(pred_vec)
            if gt_norm == 0 or pred_norm == 0:
                raise ValueError(f"Zero-norm embedding for {gt_label!r} or {pred_label!r}")
            similarity = float(np.dot(gt_vec / gt_norm, pred_vec / pred_norm))
        else:
            raise ValueError(f"Unsupported similarity mode: {self.similarity_mode}")
        return similarity >= self.similarity_threshold


def load_json(path: Path):
    with Path(path).open(encoding="utf-8") as file:
        return json.load(file)


def prepare_ground_truth(gt_payload: Mapping) -> GroundTruth:
    category_names = {
        int(category["id"]): str(category["name"])
        for category in gt_payload.get("categories", [])
        if "id" in category and "name" in category
    }
    image_ids = tuple(sorted(int(image["id"]) for image in gt_payload.get("images", [])))
    annotations_by_image = {image_id: [] for image_id in image_ids}
    for ann in gt_payload.get("annotations", []):
        image_id = int(ann["image_id"])
        if image_id in annotations_by_image:
            annotations_by_image[image_id].append(dict(ann))

    return GroundTruth(
        image_ids=image_ids,
        annotations_by_image={
            image_id: tuple(annotations_by_image[image_id]) for image_id in image_ids
        },
        category_names=category_names,
    )


def bbox_iou(left: Sequence[float], right: Sequence[float]) -> float:
    lx, ly, lw, lh = [float(value) for value in left]
    rx, ry, rw, rh = [float(value) for value in right]
    left_x2 = lx + max(lw, 0.0)
    left_y2 = ly + max(lh, 0.0)
    right_x2 = rx + max(rw, 0.0)
    right_y2 = ry + max(rh, 0.0)
    inter_w = max(0.0, min(left_x2, right_x2) - max(lx, rx))
    inter_h = max(0.0, min(left_y2, right_y2) - max(ly, ry))
    intersection = inter_w * inter_h
    union = max(lw, 0.0) * max(lh, 0.0) + max(rw, 0.0) * max(rh, 0.0) - intersection
    return intersection / union if union > 0 else 0.0


def category_label(category_id: object, category_names: Mapping[int, str]) -> str:
    if isinstance(category_id, int):
        return category_names.get(category_id, str(category_id))
    if isinstance(category_id, float) and category_id.is_integer():
        return category_names.get(int(category_id), str(int(category_id)))
    return str(category_id)


def build_category_matcher(
    dimension: str,
    category_names: Mapping[int, str],
    semantic_matcher: Optional[SemanticMatcher] = None,
) -> Callable[[object, object], bool]:
    normalized_dimension = normalize_dimension(dimension)
    if normalized_dimension == "semantics":
        if semantic_matcher is None:
            raise ValueError("--embedding-cache is required for semantic evaluation")

        def semantic_match(gt_category, pred_category):
            return semantic_matcher.matches(
                category_label(gt_category, category_names),
                category_label(pred_category, category_names),
            )

        return semantic_match

    def exact_match(gt_category, pred_category):
        return (
            SemanticMatcher.standardize(category_label(gt_category, category_names))
            == SemanticMatcher.standardize(category_label(pred_category, category_names))
        )

    return exact_match


def normalize_dimension(dimension: str) -> str:
    if dimension in {"i", "interactable", "interaction"}:
        return "interactable"
    if dimension in {"s", "semantics", "semantic"}:
        return "semantics"
    raise ValueError(f"Unsupported dimension: {dimension}")


def _score(prediction: Mapping) -> float:
    return float(prediction.get("score", 1.0))


def _empty_image_metrics(support: int, predictions: int) -> ImageMetrics:
    if support == 0 and predictions == 0:
        return ImageMetrics(precision=1.0, recall=1.0, f1=1.0, support=support, predictions=predictions)
    return ImageMetrics(precision=0.0, recall=0.0, f1=0.0, support=support, predictions=predictions)


def _metrics_from_counts(tp: int, fp: int, fn: int, support: int, predictions: int) -> ImageMetrics:
    if tp == 0 and fp == 0 and fn == 0:
        return _empty_image_metrics(support, predictions)
    precision = tp / (tp + fp) if tp + fp > 0 else 0.0
    recall = tp / (tp + fn) if tp + fn > 0 else 0.0
    f1 = 2 * precision * recall / (precision + recall) if precision + recall > 0 else 0.0
    return ImageMetrics(
        precision=precision,
        recall=recall,
        f1=f1,
        support=support,
        predictions=predictions,
    )


def group_predictions(
    predictions: Sequence[Mapping],
    image_ids: Iterable[int],
    score_threshold: float,
) -> Dict[int, List[dict]]:
    grouped = {int(image_id): [] for image_id in image_ids}
    for prediction in predictions:
        image_id = int(prediction["image_id"])
        if image_id not in grouped or _score(prediction) < score_threshold:
            continue
        grouped[image_id].append(dict(prediction))
    for image_predictions in grouped.values():
        image_predictions.sort(key=_score, reverse=True)
    return grouped


def evaluate_predictions(
    gt: GroundTruth,
    predictions: Sequence[Mapping],
    *,
    dimension: str,
    iou_threshold: float,
    score_threshold: float,
    semantic_matcher: Optional[SemanticMatcher] = None,
) -> EvaluationResult:
    match_category = build_category_matcher(dimension, gt.category_names, semantic_matcher)
    predictions_by_image = group_predictions(predictions, gt.image_ids, score_threshold)

    per_image = {}
    totals = {"tp": 0, "fp": 0, "fn": 0}
    for image_id in gt.image_ids:
        annotations = gt.annotations_by_image[image_id]
        image_predictions = predictions_by_image[image_id]
        matched_gt = set()
        tp = 0
        fp = 0
        for prediction in image_predictions:
            best_index = None
            best_iou = iou_threshold
            for ann_index, annotation in enumerate(annotations):
                if ann_index in matched_gt:
                    continue
                if not match_category(annotation["category_id"], prediction["category_id"]):
                    continue
                overlap = bbox_iou(annotation["bbox"], prediction["bbox"])
                if overlap >= best_iou:
                    best_index = ann_index
                    best_iou = overlap
            if best_index is None:
                fp += 1
            else:
                tp += 1
                matched_gt.add(best_index)
        fn = len(annotations) - len(matched_gt)
        totals["tp"] += tp
        totals["fp"] += fp
        totals["fn"] += fn
        per_image[image_id] = _metrics_from_counts(
            tp,
            fp,
            fn,
            support=len(annotations),
            predictions=len(image_predictions),
        )

    metrics = micro_metrics(totals)
    return EvaluationResult(
        metrics=metrics,
        per_image=per_image,
        counts={
            **totals,
            "gt_annotations": sum(len(anns) for anns in gt.annotations_by_image.values()),
            "predictions_before_threshold": len(predictions),
            "predictions_after_threshold": sum(
                metrics.predictions for metrics in per_image.values()
            ),
            "images": len(gt.image_ids),
        },
        score_threshold=score_threshold,
    )


def micro_metrics(counts: Mapping[str, int]) -> Dict[str, float]:
    tp = counts["tp"]
    fp = counts["fp"]
    fn = counts["fn"]
    precision = tp / (tp + fp) if tp + fp > 0 else 0.0
    recall = tp / (tp + fn) if tp + fn > 0 else 0.0
    f1 = 2 * precision * recall / (precision + recall) if precision + recall > 0 else 0.0
    return {"precision": precision, "recall": recall, "f1": f1}


def select_best_threshold(
    gt: GroundTruth,
    predictions: Sequence[Mapping],
    *,
    dimension: str,
    iou_threshold: float,
    semantic_matcher: Optional[SemanticMatcher] = None,
) -> Tuple[float, EvaluationResult]:
    best_threshold = select_best_threshold_fast(
        gt,
        predictions,
        dimension=dimension,
        iou_threshold=iou_threshold,
        semantic_matcher=semantic_matcher,
    )
    best_result = evaluate_predictions(
        gt,
        predictions,
        dimension=dimension,
        iou_threshold=iou_threshold,
        score_threshold=best_threshold,
        semantic_matcher=semantic_matcher,
    )
    return best_threshold, best_result


def select_best_threshold_fast(
    gt: GroundTruth,
    predictions: Sequence[Mapping],
    *,
    dimension: str,
    iou_threshold: float,
    semantic_matcher: Optional[SemanticMatcher] = None,
) -> float:
    if not predictions:
        return 0.0

    match_category = build_category_matcher(dimension, gt.category_names, semantic_matcher)
    predictions_by_image = group_predictions(predictions, gt.image_ids, score_threshold=-math.inf)
    total_gt = sum(len(anns) for anns in gt.annotations_by_image.values())

    events = []
    known_image_ids = set(gt.image_ids)
    for prediction in predictions:
        if int(prediction["image_id"]) not in known_image_ids:
            events.append((_score(prediction), 0, 0))

    for image_id, image_predictions in predictions_by_image.items():
        annotations = gt.annotations_by_image[image_id]
        matched_gt = set()
        for prediction in image_predictions:
            best_index = None
            best_iou = iou_threshold
            for ann_index, annotation in enumerate(annotations):
                if ann_index in matched_gt:
                    continue
                if not match_category(annotation["category_id"], prediction["category_id"]):
                    continue
                overlap = bbox_iou(annotation["bbox"], prediction["bbox"])
                if overlap >= best_iou:
                    best_index = ann_index
                    best_iou = overlap
            if best_index is None:
                events.append((_score(prediction), 0, 1))
            else:
                matched_gt.add(best_index)
                events.append((_score(prediction), 1, 0))

    if not events:
        return min(_score(prediction) for prediction in predictions)

    events.sort(key=lambda event: event[0], reverse=True)
    best_threshold = events[0][0]
    best_rank = None
    tp = 0
    fp = 0
    index = 0
    while index < len(events):
        threshold = events[index][0]
        while index < len(events) and events[index][0] == threshold:
            tp += events[index][1]
            fp += events[index][2]
            index += 1
        metrics = micro_metrics({"tp": tp, "fp": fp, "fn": total_gt - tp})
        rank = (metrics["f1"], metrics["precision"], metrics["recall"], -threshold)
        if best_rank is None or rank > best_rank:
            best_threshold = threshold
            best_rank = rank

    return best_threshold


def mean_metric(per_image: Mapping[int, ImageMetrics], image_ids: Sequence[int], metric: str) -> float:
    if not image_ids:
        return 0.0
    return statistics.fmean(getattr(per_image[image_id], metric) for image_id in image_ids)


def paired_bootstrap(
    result_a: EvaluationResult,
    result_b: EvaluationResult,
    image_ids: Sequence[int],
    *,
    replicates: int,
    seed: int,
) -> Dict[str, Mapping]:
    rng = np.random.default_rng(seed)
    deltas = {"precision": [], "recall": [], "f1": []}
    image_ids = tuple(image_ids)
    image_count = len(image_ids)
    if image_count == 0:
        return {
            metric: {
                "observed_delta": 0.0,
                "ci95_percentile": [0.0, 0.0],
                "probability_delta_gt_0": 0.0,
                "two_sided_bootstrap_p_style": 1.0,
            }
            for metric in deltas
        }
    for _ in range(replicates):
        sampled = tuple(image_ids[index] for index in rng.integers(0, image_count, size=image_count))
        for metric in deltas:
            deltas[metric].append(
                mean_metric(result_b.per_image, sampled, metric)
                - mean_metric(result_a.per_image, sampled, metric)
            )

    report = {}
    for metric, values in deltas.items():
        arr = np.asarray(values, dtype=float)
        observed_delta = mean_metric(result_b.per_image, image_ids, metric) - mean_metric(
            result_a.per_image, image_ids, metric
        )
        report[metric] = {
            "observed_delta": observed_delta,
            "ci95_percentile": [float(np.percentile(arr, 2.5)), float(np.percentile(arr, 97.5))],
            "probability_delta_gt_0": float(np.mean(arr > 0)),
            "two_sided_bootstrap_p_style": float(
                min(1.0, 2 * min(np.mean(arr <= 0), np.mean(arr >= 0)))
            ),
        }
    return report


def serialize_evaluation(result: EvaluationResult) -> Dict[str, object]:
    all_images = tuple(result.per_image.keys())
    positive_support_images = tuple(
        image_id for image_id, metrics in result.per_image.items() if metrics.support > 0
    )
    all_image_mean = {
        metric: mean_metric(result.per_image, all_images, metric)
        for metric in ("precision", "recall", "f1")
    }
    positive_support_mean = {
        metric: mean_metric(result.per_image, positive_support_images, metric)
        for metric in ("precision", "recall", "f1")
    }
    return {
        "score_threshold": result.score_threshold,
        "point": {
            "micro": dict(result.metrics),
            "mean_per_all_images": all_image_mean,
            "mean_per_positive_support_images": positive_support_mean,
        },
        "counts": dict(result.counts),
    }


def build_report(
    gt: GroundTruth,
    predictions_a: Sequence[Mapping],
    predictions_b: Sequence[Mapping],
    *,
    dimension: str,
    iou_threshold: float,
    score_threshold: float,
    auto_threshold: bool,
    embedding_cache: Optional[Path],
    replicates: int,
    seed: int,
) -> Dict[str, object]:
    normalized_dimension = normalize_dimension(dimension)
    semantic_matcher = (
        SemanticMatcher(embedding_cache)
        if normalized_dimension == "semantics" and embedding_cache is not None
        else None
    )

    if auto_threshold:
        _, result_a = select_best_threshold(
            gt,
            predictions_a,
            dimension=normalized_dimension,
            iou_threshold=iou_threshold,
            semantic_matcher=semantic_matcher,
        )
        _, result_b = select_best_threshold(
            gt,
            predictions_b,
            dimension=normalized_dimension,
            iou_threshold=iou_threshold,
            semantic_matcher=semantic_matcher,
        )
    else:
        result_a = evaluate_predictions(
            gt,
            predictions_a,
            dimension=normalized_dimension,
            iou_threshold=iou_threshold,
            score_threshold=score_threshold,
            semantic_matcher=semantic_matcher,
        )
        result_b = evaluate_predictions(
            gt,
            predictions_b,
            dimension=normalized_dimension,
            iou_threshold=iou_threshold,
            score_threshold=score_threshold,
            semantic_matcher=semantic_matcher,
        )

    return {
        "protocol": {
            "dimension": normalized_dimension,
            "iou_threshold": iou_threshold,
            "auto_threshold": auto_threshold,
            "fixed_score_threshold": None if auto_threshold else score_threshold,
            "semantic_similarity_threshold": (
                SIMILARITY_THRESHOLD if normalized_dimension == "semantics" else None
            ),
            "semantic_similarity_mode": (
                "raw_dot_product_historical" if normalized_dimension == "semantics" else None
            ),
            "bootstrap_unit": "image_id",
            "replicates": replicates,
            "seed": seed,
            "empty_image_convention": (
                "Images with no GT annotations and no retained predictions receive per-image "
                "precision=recall=F1=1.0 for all-image absence accounting; positive-support "
                "means exclude those images."
            ),
            "caveat": (
                "This paired image-cluster bootstrap complements the paper's COCO-style "
                "macro-category max-F1 evaluation; it is not a byte-for-byte duplicate "
                "of COCOeval aggregation. When auto-threshold is enabled, thresholds are "
                "selected on the same ground truth and held fixed during resampling, so the "
                "intervals are conditional, exploratory, and may be optimistic."
            ),
        },
        "coverage": {
            "images": len(gt.image_ids),
            "gt_annotations": sum(len(anns) for anns in gt.annotations_by_image.values()),
            "gt_categories": len(gt.category_names),
            "embedding_cache_entries": (
                len(semantic_matcher.embeddings) if semantic_matcher is not None else None
            ),
            "embedding_dimension": (
                semantic_matcher.dimension if semantic_matcher is not None else None
            ),
        },
        "methods": {
            "method_a": serialize_evaluation(result_a),
            "method_b": serialize_evaluation(result_b),
        },
        "bootstrap": {
            "delta_direction": "method_b_minus_method_a",
            "mean_per_all_images_delta": paired_bootstrap(
                result_a,
                result_b,
                gt.image_ids,
                replicates=replicates,
                seed=seed,
            ),
            "mean_per_positive_support_images_delta": paired_bootstrap(
                result_a,
                result_b,
                tuple(
                    image_id
                    for image_id in gt.image_ids
                    if len(gt.annotations_by_image[image_id]) > 0
                ),
                replicates=replicates,
                seed=seed,
            ),
        },
    }


def build_parser():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--gt", type=Path, required=True, help="COCO ground-truth JSON")
    parser.add_argument("--method-a", type=Path, required=True, help="first COCO result JSON")
    parser.add_argument("--method-b", type=Path, required=True, help="second COCO result JSON")
    parser.add_argument(
        "--dimension",
        choices=("interactable", "interaction", "i", "semantics", "semantic", "s"),
        required=True,
    )
    parser.add_argument("--embedding-cache", type=Path, help="required for semantic mode")
    parser.add_argument("--iou-threshold", type=float, default=0.75)
    parser.add_argument("--score-threshold", type=float, default=0.0)
    parser.add_argument("--auto-threshold", action="store_true")
    parser.add_argument("--replicates", type=int, default=10000)
    parser.add_argument("--seed", type=int, default=20260822)
    parser.add_argument("--output", type=Path, help="write JSON report to this path")
    return parser


def main(argv=None):
    args = build_parser().parse_args(argv)
    if args.replicates <= 0:
        raise ValueError("--replicates must be positive")
    normalized_dimension = normalize_dimension(args.dimension)
    if normalized_dimension == "semantics" and args.embedding_cache is None:
        raise ValueError("--embedding-cache is required for semantic evaluation")

    report = build_report(
        prepare_ground_truth(load_json(args.gt)),
        load_json(args.method_a),
        load_json(args.method_b),
        dimension=normalized_dimension,
        iou_threshold=args.iou_threshold,
        score_threshold=args.score_threshold,
        auto_threshold=args.auto_threshold,
        embedding_cache=args.embedding_cache,
        replicates=args.replicates,
        seed=args.seed,
    )
    output = json.dumps(report, indent=2, sort_keys=True)
    if args.output:
        args.output.write_text(output + "\n", encoding="utf-8")
    else:
        print(output)


if __name__ == "__main__":
    main()