from __future__ import annotations

import argparse
import csv
import json
import math
import random
import sys
import time
from collections import Counter, defaultdict
from datetime import datetime
from pathlib import Path
from typing import Any, Optional

import cv2
import numpy as np
import torch

FILE = Path(__file__).resolve()
ROOT = FILE.parents[2]
if str(ROOT) not in sys.path:
    sys.path.append(str(ROOT))

from tools.pdcl_inference.analyze_val_two_roi_badcases import (
    MIN_CONFIDENCE_FOR_2D_THRESHOLD_SEARCH,
    annotate_panel_title,
    box_iou_matrix,
    draw_box_with_label,
    entry_to_image_file,
    entry_to_label_file,
    get_cls_name,
    greedy_match_indices_any_class,
    infer_class_name_map,
    load_split_entries,
    load_yaml,
    make_reservoir_store,
    names_to_dict,
    prepare_gt_for_roi,
    read_raw_calib_from_label,
    reservoir_add,
    reservoir_records,
    resolve_data_path,
    run_model_for_prepared_roi,
    sanitize_name,
    to_float,
)
from tools.pdcl_inference.two_roi_inference import (
    _filter_prediction_rows,
    _prepare_roi_image,
    build_inference_context,
)
from tools.pdcl_inference.run_batch_two_roi_infer import (
    add_two_roi_inference_args,
    build_roi_specs_from_args,
)


DEFAULT_OUTPUT_ROOT = FILE.parent / "validation_analysis" / "background_miss_samples_{}".format(
    datetime.now().strftime("%Y%m%d_%H%M%S")
)
MISS_REASON_ORDER = ("no_overlap", "localization", "threshold_limited", "assignment_conflict")
MISS_RECORD_FIELDS = [
    "sample_index",
    "roi",
    "frame_name",
    "image_path",
    "label_path",
    "cls_id",
    "cls_name",
    "gt_index",
    "confidence_threshold",
    "threshold_source",
    "miss_reason",
    "difficulty",
    "gt_bbox_xyxy",
    "gt_bbox_w_px",
    "gt_bbox_h_px",
    "gt_bbox_diag_px",
    "max_iou_kept",
    "best_kept_pred_index",
    "best_kept_pred_cls_id",
    "best_kept_pred_cls_name",
    "best_kept_pred_conf",
    "best_kept_pred_bbox_xyxy",
    "max_iou_raw",
    "best_raw_pred_index",
    "best_raw_pred_cls_id",
    "best_raw_pred_cls_name",
    "best_raw_pred_conf",
    "best_raw_pred_bbox_xyxy",
    "visualization",
]


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Sample GT objects that end up in the confusion-matrix background row (class-agnostic 2D misses)."
    )
    parser.add_argument(
        "--data",
        type=str,
        default=str(ROOT / "ultralytics" / "cfg" / "datasets" / "mono3d_ground.yaml"),
        help="Dataset YAML path used to resolve the analyzed split.",
    )
    parser.add_argument("--split", type=str, default="val", choices=("train", "val", "test"), help="Dataset split to analyze.")
    parser.add_argument(
        "--output-root",
        type=str,
        default=str(DEFAULT_OUTPUT_ROOT),
        help="Directory where sampled miss visualizations and manifests will be written.",
    )
    parser.add_argument(
        "--analyze-rois",
        nargs="+",
        default=["roi0", "roi1"],
        choices=("roi0", "roi1"),
        help="ROI models to analyze.",
    )
    parser.add_argument("--max-samples", type=int, default=0, help="Optional cap on analyzed samples. 0 means full split.")
    parser.add_argument(
        "--sample-selection",
        type=str,
        default="head",
        choices=("head", "random"),
        help="How to choose the first --max-samples entries when a cap is set.",
    )
    parser.add_argument("--sample-random-seed", type=int, default=20260408, help="Random seed used when --sample-selection=random.")
    parser.add_argument("--per-class-samples", type=int, default=100, help="How many random missed GT samples to save per class.")
    parser.add_argument("--badcase-random-seed", type=int, default=20260408, help="Random seed for per-class reservoir sampling.")
    parser.add_argument("--log-every", type=int, default=100, help="Progress log interval in samples.")
    parser.add_argument("--torch-threads", type=int, default=0, help="Optional torch CPU thread count override.")
    parser.add_argument(
        "--confidence-threshold",
        type=float,
        default=None,
        help="Explicit confidence threshold for the confusion-matrix miss definition. If unset, try --report-root first, else use the model default.",
    )
    parser.add_argument(
        "--report-root",
        type=str,
        default=None,
        help="Optional existing report root whose <roi>/summary.json provides recommended_2d_confidence.",
    )
    parser.add_argument(
        "--raw-conf-threshold",
        type=float,
        default=float(MIN_CONFIDENCE_FOR_2D_THRESHOLD_SEARCH),
        help="Low confidence threshold used to keep raw candidate detections for debugging threshold-limited misses.",
    )
    parser.add_argument(
        "--crop-scale",
        type=float,
        default=2.5,
        help="Context scale around the GT box for the zoomed panel.",
    )
    add_two_roi_inference_args(parser, include_output_dir=False)
    return parser.parse_args()


def load_json(path: Path) -> dict[str, Any]:
    with path.open("r", encoding="utf-8") as file:
        return json.load(file)


def write_json(path: Path, payload: Any) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8") as file:
        json.dump(payload, file, indent=2, ensure_ascii=False)


def write_csv(path: Path, rows: list[dict[str, Any]], fieldnames: list[str]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8", newline="") as file:
        writer = csv.DictWriter(file, fieldnames=fieldnames)
        writer.writeheader()
        for row in rows:
            writer.writerow(row)


def maybe_list(value: Any) -> Optional[list[float]]:
    if value is None:
        return None
    return [float(v) for v in np.asarray(value, dtype=np.float32).reshape(-1).tolist()]


def float_or_none(value: Any) -> Optional[float]:
    resolved = to_float(value)
    return None if resolved is None or not math.isfinite(float(resolved)) else float(resolved)


def resolve_confidence_threshold(
    roi_name: str,
    bundle,
    explicit_threshold: Optional[float],
    report_root: Optional[Path],
) -> tuple[float, str]:
    if explicit_threshold is not None:
        return float(explicit_threshold), "explicit_arg"
    if report_root is not None:
        summary_path = report_root / roi_name / "summary.json"
        if summary_path.is_file():
            payload = load_json(summary_path)
            overall = payload.get("overall") or {}
            threshold_advice = payload.get("threshold_advice_2d") or {}
            recommended = float_or_none(overall.get("recommended_confidence_2d"))
            if recommended is None:
                recommended = float_or_none(threshold_advice.get("recommended_confidence"))
            if recommended is not None:
                return float(recommended), f"report:{summary_path}"
    return float(bundle.spec.conf), "model_default"


def classify_miss_reason(max_iou_kept: Optional[float], max_iou_raw: Optional[float]) -> str:
    kept = 0.0 if max_iou_kept is None else float(max_iou_kept)
    raw = 0.0 if max_iou_raw is None else float(max_iou_raw)
    if kept >= 0.5:
        return "assignment_conflict"
    if raw >= 0.5:
        return "threshold_limited"
    if raw >= 0.1:
        return "localization"
    return "no_overlap"


def crop_with_context(image: np.ndarray, gt_box: list[float], crop_scale: float) -> tuple[np.ndarray, int, int]:
    x1, y1, x2, y2 = [float(v) for v in gt_box]
    w = max(x2 - x1, 1.0)
    h = max(y2 - y1, 1.0)
    cx = 0.5 * (x1 + x2)
    cy = 0.5 * (y1 + y2)
    side = max(w, h) * max(float(crop_scale), 1.0)
    side = max(side, 160.0)
    crop_x1 = max(0, int(math.floor(cx - side * 0.5)))
    crop_y1 = max(0, int(math.floor(cy - side * 0.5)))
    crop_x2 = min(image.shape[1], int(math.ceil(cx + side * 0.5)))
    crop_y2 = min(image.shape[0], int(math.ceil(cy + side * 0.5)))
    if crop_x2 <= crop_x1:
        crop_x2 = min(image.shape[1], crop_x1 + 1)
    if crop_y2 <= crop_y1:
        crop_y2 = min(image.shape[0], crop_y1 + 1)
    return image[crop_y1:crop_y2, crop_x1:crop_x2].copy(), crop_x1, crop_y1


def translate_box(xyxy: Optional[list[float]], offset_x: int, offset_y: int) -> Optional[list[float]]:
    if xyxy is None:
        return None
    x1, y1, x2, y2 = [float(v) for v in xyxy]
    return [x1 - float(offset_x), y1 - float(offset_y), x2 - float(offset_x), y2 - float(offset_y)]


def draw_record_overlays(image: np.ndarray, record: dict[str, Any], title: str) -> np.ndarray:
    panel = draw_box_with_label(image, record["gt_bbox_xyxy"], (0, 200, 0), f"GT {record['cls_name']}")
    kept_box = record.get("best_kept_pred_bbox_xyxy")
    if kept_box is not None:
        kept_label = (
            f"kept {record.get('best_kept_pred_cls_name', 'unknown')} "
            f"{record.get('best_kept_pred_conf', 0.0):.2f} IoU={record.get('max_iou_kept', 0.0):.2f}"
        )
        panel = draw_box_with_label(panel, kept_box, (0, 0, 255), kept_label)
    raw_box = record.get("best_raw_pred_bbox_xyxy")
    raw_index = record.get("best_raw_pred_index")
    kept_index = record.get("best_kept_pred_index")
    if raw_box is not None and raw_index != kept_index:
        raw_label = (
            f"raw {record.get('best_raw_pred_cls_name', 'unknown')} "
            f"{record.get('best_raw_pred_conf', 0.0):.2f} IoU={record.get('max_iou_raw', 0.0):.2f}"
        )
        panel = draw_box_with_label(panel, raw_box, (0, 165, 255), raw_label)
    return annotate_panel_title(panel, title)


def make_text_panel(shape: tuple[int, int, int], title: str, record: dict[str, Any]) -> np.ndarray:
    panel = np.full(shape, 28, dtype=np.uint8)
    lines = [
        title,
        f"{record['cls_name']} miss_reason={record['miss_reason']}",
        f"threshold={record['confidence_threshold']:.3f} source={record['threshold_source']}",
        (
            f"GT diag={record['gt_bbox_diag_px']:.1f}px "
            f"w={record['gt_bbox_w_px']:.1f}px h={record['gt_bbox_h_px']:.1f}px"
        ),
        f"difficulty={record.get('difficulty') if record.get('difficulty') is not None else 'n/a'}",
        (
            f"kept IoU={record.get('max_iou_kept', 0.0):.3f} "
            f"cls={record.get('best_kept_pred_cls_name', 'n/a')} conf={record.get('best_kept_pred_conf', 'n/a')}"
        ),
        (
            f"raw IoU={record.get('max_iou_raw', 0.0):.3f} "
            f"cls={record.get('best_raw_pred_cls_name', 'n/a')} conf={record.get('best_raw_pred_conf', 'n/a')}"
        ),
        f"frame={record['frame_name']}",
        f"gt_index={record['gt_index']} sample_index={record['sample_index']}",
    ]
    y = 28
    for line in lines:
        cv2.putText(panel, str(line), (10, y), cv2.FONT_HERSHEY_SIMPLEX, 0.56, (220, 220, 220), 1, cv2.LINE_AA)
        y += 34
    return panel


def save_sample_visuals(
    records: list[dict[str, Any]],
    output_dir: Path,
    bundle,
    crop_scale: float,
) -> list[dict[str, Any]]:
    output_dir.mkdir(parents=True, exist_ok=True)
    manifest: list[dict[str, Any]] = []
    for rank, record in enumerate(records, start=1):
        image_path = Path(str(record["image_path"]))
        label_path = Path(str(record["label_path"]))
        image_bgr = cv2.imread(str(image_path), cv2.IMREAD_COLOR)
        if image_bgr is None:
            continue
        raw_calib = read_raw_calib_from_label(image_path, label_path)
        prepared = _prepare_roi_image(image_bgr, raw_calib, bundle.spec, bundle.imgsz)
        roi_image = prepared.image.copy()
        panel_size = (roi_image.shape[1], roi_image.shape[0])

        panel_full = draw_record_overlays(roi_image, record, f"{bundle.spec.name} full miss")
        crop_image, crop_x1, crop_y1 = crop_with_context(roi_image, record["gt_bbox_xyxy"], crop_scale)
        crop_record = dict(record)
        crop_record["gt_bbox_xyxy"] = translate_box(record["gt_bbox_xyxy"], crop_x1, crop_y1)
        crop_record["best_kept_pred_bbox_xyxy"] = translate_box(record.get("best_kept_pred_bbox_xyxy"), crop_x1, crop_y1)
        crop_record["best_raw_pred_bbox_xyxy"] = translate_box(record.get("best_raw_pred_bbox_xyxy"), crop_x1, crop_y1)
        panel_crop = draw_record_overlays(crop_image, crop_record, f"{bundle.spec.name} crop")
        panel_crop = cv2.resize(panel_crop, panel_size, interpolation=cv2.INTER_LINEAR)
        panel_text = make_text_panel(roi_image.shape, f"{bundle.spec.name} miss #{rank}", record)

        grid = np.concatenate([panel_full, panel_crop, panel_text], axis=1)
        filename = (
            f"{rank:03d}_{sanitize_name(Path(record['frame_name']).stem)}_{sanitize_name(record['cls_name'])}"
            f"_reason_{sanitize_name(record['miss_reason'])}_g{record['gt_index']}.jpg"
        )
        image_out = output_dir / filename
        cv2.imwrite(str(image_out), grid)
        manifest.append({**record, "visualization": str(image_out)})

    write_json(output_dir / "manifest.json", manifest)
    return manifest


def record_to_csv_row(record: dict[str, Any]) -> dict[str, Any]:
    row: dict[str, Any] = {}
    for field in MISS_RECORD_FIELDS:
        value = record.get(field)
        if isinstance(value, (list, tuple, dict)):
            row[field] = json.dumps(value, ensure_ascii=False)
        else:
            row[field] = value
    return row


def sample_background_misses_for_roi(
    bundle,
    entries: list[tuple[str, str]],
    image_root: Path,
    class_map: dict[str, int],
    difficulty_weights: list[float],
    face_3d_classes: set[int],
    complete_3d_classes: set[int],
    classes: Optional[set[int]],
    min_wh_px: float,
    output_root: Path,
    per_class_samples: int,
    log_every: int,
    badcase_random_seed: int,
    confidence_threshold: float,
    threshold_source: str,
    raw_conf_threshold: float,
    crop_scale: float,
) -> dict[str, Any]:
    roi_name = bundle.spec.name.lower()
    roi_output = output_root / roi_name
    roi_output.mkdir(parents=True, exist_ok=True)
    names_dict = names_to_dict(bundle.names)

    class_gt_total: Counter[int] = Counter()
    class_miss_total: Counter[int] = Counter()
    class_reason_counts: defaultdict[int, Counter[str]] = defaultdict(Counter)
    class_sample_stores: defaultdict[tuple[int, str], dict[str, Any]] = defaultdict(make_reservoir_store)
    rng = random.Random(int(badcase_random_seed) + (0 if roi_name == "roi0" else 1000))
    search_conf = min(float(raw_conf_threshold), float(bundle.spec.conf))
    start_time = time.time()

    for sample_index, entry in enumerate(entries):
        image_path = entry_to_image_file(entry, image_root)
        label_path = entry_to_label_file(entry)
        image_bgr = cv2.imread(str(image_path), cv2.IMREAD_COLOR)
        if image_bgr is None:
            continue
        ori_h, ori_w = image_bgr.shape[:2]
        raw_calib = read_raw_calib_from_label(image_path, label_path)
        prepared = _prepare_roi_image(image_bgr, raw_calib, bundle.spec, bundle.imgsz)
        gt = prepare_gt_for_roi(
            label_file=label_path,
            ori_w=ori_w,
            ori_h=ori_h,
            prepared=prepared,
            bundle=bundle,
            class_map=class_map,
            difficulty_weights=difficulty_weights,
            face_3d_classes=face_3d_classes,
            complete_3d_classes=complete_3d_classes,
            include_classes=classes,
            min_wh_px=min_wh_px,
        )
        raw_outputs = run_model_for_prepared_roi(bundle, prepared)
        analysis_detections, _analysis_preds_3d, _analysis_preds_edge, _analysis_anchors, _analysis_strides = _filter_prediction_rows(
            *raw_outputs,
            conf_thres=float(search_conf),
            max_det=bundle.spec.max_det,
            classes=classes,
        )
        raw_pred_boxes = (
            np.asarray(analysis_detections[:, :4], dtype=np.float32) if len(analysis_detections) else np.zeros((0, 4), dtype=np.float32)
        )
        raw_pred_cls = (
            np.asarray(analysis_detections[:, 5], dtype=np.int32).reshape(-1) if len(analysis_detections) else np.zeros((0,), dtype=np.int32)
        )
        raw_pred_conf = (
            np.asarray(analysis_detections[:, 4], dtype=np.float32).reshape(-1) if len(analysis_detections) else np.zeros((0,), dtype=np.float32)
        )
        keep = raw_pred_conf > float(confidence_threshold) if raw_pred_conf.size else np.zeros((0,), dtype=bool)
        kept_pred_boxes = raw_pred_boxes[keep]
        kept_pred_cls = raw_pred_cls[keep]
        kept_pred_conf = raw_pred_conf[keep]

        gt_boxes = np.asarray(gt["boxes_xyxy"], dtype=np.float32).reshape(-1, 4)
        gt_cls = np.asarray(gt["classes"], dtype=np.int32).reshape(-1)
        gt_difficulties = np.asarray(gt["lb_2d"].get("difficulties", []), dtype=np.float32).reshape(-1)
        matches, iou_kept = greedy_match_indices_any_class(gt_boxes, kept_pred_boxes, iou_thr=0.5)
        iou_raw = box_iou_matrix(gt_boxes, raw_pred_boxes)
        matched_gt = {int(gt_index) for gt_index, _ in matches.tolist()}

        for cls_id in gt_cls.tolist():
            class_gt_total[int(cls_id)] += 1

        for gt_index, cls_id in enumerate(gt_cls.tolist()):
            if int(gt_index) in matched_gt:
                continue
            cls_id = int(cls_id)
            cls_name = get_cls_name(names_dict, cls_id)
            class_miss_total[cls_id] += 1

            kept_row = iou_kept[gt_index] if iou_kept.shape[1] > 0 else np.zeros((0,), dtype=np.float32)
            raw_row = iou_raw[gt_index] if iou_raw.shape[1] > 0 else np.zeros((0,), dtype=np.float32)
            best_kept_pred_index = None if kept_row.size == 0 else int(np.argmax(kept_row))
            best_raw_pred_index = None if raw_row.size == 0 else int(np.argmax(raw_row))
            max_iou_kept = 0.0 if best_kept_pred_index is None else float(kept_row[best_kept_pred_index])
            max_iou_raw = 0.0 if best_raw_pred_index is None else float(raw_row[best_raw_pred_index])
            if best_kept_pred_index is not None and max_iou_kept <= 0.0:
                best_kept_pred_index = None
                max_iou_kept = 0.0
            if best_raw_pred_index is not None and max_iou_raw <= 0.0:
                best_raw_pred_index = None
                max_iou_raw = 0.0
            miss_reason = classify_miss_reason(max_iou_kept=max_iou_kept, max_iou_raw=max_iou_raw)
            class_reason_counts[cls_id][miss_reason] += 1

            gt_box = gt_boxes[gt_index]
            gt_w = float(max(gt_box[2] - gt_box[0], 0.0))
            gt_h = float(max(gt_box[3] - gt_box[1], 0.0))
            record = {
                "sample_index": int(sample_index),
                "roi": roi_name,
                "frame_name": image_path.name,
                "image_path": str(image_path),
                "label_path": str(label_path),
                "cls_id": cls_id,
                "cls_name": cls_name,
                "gt_index": int(gt_index),
                "confidence_threshold": float(confidence_threshold),
                "threshold_source": str(threshold_source),
                "miss_reason": miss_reason,
                "difficulty": (
                    None if gt_index >= len(gt_difficulties) else int(round(float(gt_difficulties[gt_index])))
                ),
                "gt_bbox_xyxy": maybe_list(gt_box),
                "gt_bbox_w_px": gt_w,
                "gt_bbox_h_px": gt_h,
                "gt_bbox_diag_px": float(math.hypot(gt_w, gt_h)),
                "max_iou_kept": float(max_iou_kept),
                "best_kept_pred_index": None if best_kept_pred_index is None else int(best_kept_pred_index),
                "best_kept_pred_cls_id": (
                    None if best_kept_pred_index is None else int(kept_pred_cls[best_kept_pred_index])
                ),
                "best_kept_pred_cls_name": (
                    None if best_kept_pred_index is None else get_cls_name(names_dict, int(kept_pred_cls[best_kept_pred_index]))
                ),
                "best_kept_pred_conf": (
                    None if best_kept_pred_index is None else float(kept_pred_conf[best_kept_pred_index])
                ),
                "best_kept_pred_bbox_xyxy": (
                    None if best_kept_pred_index is None else maybe_list(kept_pred_boxes[best_kept_pred_index])
                ),
                "max_iou_raw": float(max_iou_raw),
                "best_raw_pred_index": None if best_raw_pred_index is None else int(best_raw_pred_index),
                "best_raw_pred_cls_id": None if best_raw_pred_index is None else int(raw_pred_cls[best_raw_pred_index]),
                "best_raw_pred_cls_name": (
                    None if best_raw_pred_index is None else get_cls_name(names_dict, int(raw_pred_cls[best_raw_pred_index]))
                ),
                "best_raw_pred_conf": None if best_raw_pred_index is None else float(raw_pred_conf[best_raw_pred_index]),
                "best_raw_pred_bbox_xyxy": None if best_raw_pred_index is None else maybe_list(raw_pred_boxes[best_raw_pred_index]),
            }
            reservoir_add(class_sample_stores[(cls_id, cls_name)], record, per_class_samples, rng)

        if (sample_index + 1) % max(1, log_every) == 0 or (sample_index + 1) == len(entries):
            elapsed = time.time() - start_time
            per_sample = elapsed / max(sample_index + 1, 1)
            remaining = len(entries) - (sample_index + 1)
            eta = remaining * per_sample
            print(
                f"[{roi_name}] {sample_index + 1}/{len(entries)} "
                f"elapsed={elapsed / 60:.1f}m eta={eta / 60:.1f}m background_misses={sum(class_miss_total.values())}"
            )

    class_summary_rows: list[dict[str, Any]] = []
    sampled_rows: list[dict[str, Any]] = []
    class_manifests: list[dict[str, Any]] = []
    for (cls_id, cls_name), store in sorted(class_sample_stores.items(), key=lambda item: item[0][0]):
        sampled_records = reservoir_records(store, rng)
        class_dir = roi_output / "samples_by_class" / f"{int(cls_id):02d}_{sanitize_name(cls_name)}"
        manifest = save_sample_visuals(sampled_records, class_dir, bundle=bundle, crop_scale=float(crop_scale))
        class_manifests.append(
            {
                "cls_id": int(cls_id),
                "cls_name": str(cls_name),
                "count": len(manifest),
                "manifest_path": str(class_dir / "manifest.json"),
            }
        )
        sampled_rows.extend(record_to_csv_row(record) for record in manifest)
        reason_counts = class_reason_counts.get(int(cls_id), Counter())
        class_summary_rows.append(
            {
                "cls_id": int(cls_id),
                "cls_name": str(cls_name),
                "gt_total": int(class_gt_total.get(int(cls_id), 0)),
                "background_miss_total": int(class_miss_total.get(int(cls_id), 0)),
                "background_miss_rate": (
                    float(class_miss_total.get(int(cls_id), 0)) / float(class_gt_total.get(int(cls_id), 0))
                    if int(class_gt_total.get(int(cls_id), 0)) > 0
                    else None
                ),
                "sampled_count": int(len(manifest)),
                "confidence_threshold": float(confidence_threshold),
                "threshold_source": str(threshold_source),
                "no_overlap_count": int(reason_counts.get("no_overlap", 0)),
                "localization_count": int(reason_counts.get("localization", 0)),
                "threshold_limited_count": int(reason_counts.get("threshold_limited", 0)),
                "assignment_conflict_count": int(reason_counts.get("assignment_conflict", 0)),
            }
        )

    summary_fields = [
        "cls_id",
        "cls_name",
        "gt_total",
        "background_miss_total",
        "background_miss_rate",
        "sampled_count",
        "confidence_threshold",
        "threshold_source",
        "no_overlap_count",
        "localization_count",
        "threshold_limited_count",
        "assignment_conflict_count",
    ]
    write_csv(roi_output / "class_background_miss_summary.csv", class_summary_rows, summary_fields)
    write_csv(roi_output / "sampled_background_misses.csv", sampled_rows, MISS_RECORD_FIELDS)

    payload = {
        "roi": roi_name,
        "model_path": bundle.spec.model_path,
        "confidence_threshold": float(confidence_threshold),
        "threshold_source": str(threshold_source),
        "raw_conf_threshold": float(search_conf),
        "per_class_samples": int(per_class_samples),
        "total_gt": int(sum(class_gt_total.values())),
        "total_background_misses": int(sum(class_miss_total.values())),
        "class_summary_rows": class_summary_rows,
        "class_manifests": class_manifests,
        "summary_csv": str(roi_output / "class_background_miss_summary.csv"),
        "sampled_csv": str(roi_output / "sampled_background_misses.csv"),
    }
    write_json(roi_output / "summary.json", payload)
    return payload


def main() -> None:
    args = parse_args()
    if args.torch_threads and args.torch_threads > 0:
        torch.set_num_threads(int(args.torch_threads))
        torch.set_num_interop_threads(max(1, min(int(args.torch_threads), 4)))

    data_yaml = Path(args.data).resolve()
    data_cfg = load_yaml(data_yaml)
    dataset_root = data_cfg.get("path")
    split_path = resolve_data_path(data_yaml, dataset_root, data_cfg.get(args.split))
    entries = load_split_entries(
        split_path,
        max_samples=args.max_samples,
        sample_selection=str(args.sample_selection),
        sample_random_seed=int(args.sample_random_seed),
    )
    image_root = resolve_data_path(data_yaml, None, dataset_root)

    args.roi0_data = args.roi0_data or args.data
    args.roi1_data = args.roi1_data or args.data
    requested_rois = {roi.lower() for roi in args.analyze_rois}
    roi_specs = [spec for spec in build_roi_specs_from_args(args) if spec.name.lower() in requested_rois]
    context = build_inference_context(
        roi_specs=roi_specs,
        device=args.device,
        half=args.half,
        classes=args.classes,
        edge_yaw_max_lateral_dist_m=float(args.edge_yaw_max_lateral_dist),
    )

    output_root = Path(args.output_root).resolve()
    output_root.mkdir(parents=True, exist_ok=True)
    report_root = None if args.report_root is None else Path(args.report_root).resolve()

    class_map = {str(key): int(value) for key, value in (data_cfg.get("class_map") or {}).items()}
    difficulty_weights = [float(v) for v in data_cfg.get("difficulty_weights", [1.0, 1.0, 1.0, 1.0])]
    face_3d_classes = set(int(v) for v in data_cfg.get("face_3d_classes", []))
    complete_3d_classes = set(int(v) for v in data_cfg.get("complete_3d_classes", []))
    include_classes = None if args.classes is None else set(int(v) for v in args.classes)
    min_wh_px = float(data_cfg.get("min_wh", 2.0))
    class_names = infer_class_name_map(class_map, names_to_dict(context.roi_models[0].names) if context.roi_models else None)

    start = time.time()
    summary_by_roi: dict[str, dict[str, Any]] = {}
    for bundle in context.roi_models:
        roi_name = bundle.spec.name.lower()
        confidence_threshold, threshold_source = resolve_confidence_threshold(
            roi_name=roi_name,
            bundle=bundle,
            explicit_threshold=float_or_none(args.confidence_threshold),
            report_root=report_root,
        )
        print(
            f"\nSampling confusion-matrix background misses for {roi_name} on {len(entries)} {args.split} samples "
            f"(conf>{confidence_threshold:.3f}, source={threshold_source})..."
        )
        summary_by_roi[roi_name] = sample_background_misses_for_roi(
            bundle=bundle,
            entries=entries,
            image_root=image_root,
            class_map=class_map,
            difficulty_weights=difficulty_weights,
            face_3d_classes=face_3d_classes,
            complete_3d_classes=complete_3d_classes,
            classes=include_classes,
            min_wh_px=min_wh_px,
            output_root=output_root,
            per_class_samples=int(args.per_class_samples),
            log_every=int(args.log_every),
            badcase_random_seed=int(args.badcase_random_seed),
            confidence_threshold=float(confidence_threshold),
            threshold_source=threshold_source,
            raw_conf_threshold=float(args.raw_conf_threshold),
            crop_scale=float(args.crop_scale),
        )

    combined_summary = {
        "data": str(data_yaml),
        "split": str(args.split),
        "split_path": str(split_path),
        "image_root": str(image_root),
        "num_entries": int(len(entries)),
        "output_root": str(output_root),
        "class_names": class_names,
        "summary_by_roi": summary_by_roi,
        "args": vars(args),
        "elapsed_minutes": (time.time() - start) / 60.0,
    }
    write_json(output_root / "summary.json", combined_summary)

    print("\nBackground-miss sampling complete.")
    print(f"Output root: {output_root}")
    for roi_name, payload in summary_by_roi.items():
        print(
            f"[{roi_name}] total_gt={payload['total_gt']} "
            f"background_misses={payload['total_background_misses']} "
            f"summary={payload['summary_csv']}"
        )


if __name__ == "__main__":
    main()
