edgefirst_client/coco/

verify.rs

// SPDX-License-Identifier: Apache-2.0
// Copyright © 2025 Au-Zone Technologies. All Rights Reserved.

//! COCO dataset verification utilities.
//!
//! Provides functions for comparing COCO datasets and validating annotation
//! accuracy using Hungarian matching for optimal annotation pairing.

use super::{
    decode_compressed_rle, decode_rle,
    types::{CocoAnnotation, CocoDataset, CocoSegmentation},
};
use pathfinding::{kuhn_munkres::kuhn_munkres_min, matrix::Matrix};
use std::{
    collections::{HashMap, HashSet},
    fmt,
};

/// Result of verifying a COCO import against Studio data.
#[derive(Debug, Clone)]
pub struct VerificationResult {
    /// Total images in COCO dataset.
    pub coco_image_count: usize,
    /// Images found in Studio.
    pub studio_image_count: usize,
    /// Images missing from Studio.
    pub missing_images: Vec<String>,
    /// Extra images in Studio not in COCO.
    pub extra_images: Vec<String>,
    /// Total annotations in COCO dataset.
    pub coco_annotation_count: usize,
    /// Total annotations in Studio.
    pub studio_annotation_count: usize,
    /// Bounding box validation results.
    pub bbox_validation: BboxValidationResult,
    /// Segmentation mask validation results.
    pub mask_validation: MaskValidationResult,
    /// Category validation results.
    pub category_validation: CategoryValidationResult,
}

impl VerificationResult {
    /// Returns true if the verification passed all checks.
    pub fn is_valid(&self) -> bool {
        self.missing_images.is_empty()
            && self.extra_images.is_empty()
            && self.bbox_validation.is_valid()
            && self.mask_validation.is_valid()
    }

    /// Returns a summary of the verification.
    pub fn summary(&self) -> String {
        let mut s = String::new();
        s.push_str(&format!(
            "Images: {}/{} (missing: {}, extra: {})\n",
            self.studio_image_count,
            self.coco_image_count,
            self.missing_images.len(),
            self.extra_images.len()
        ));
        s.push_str(&format!(
            "Annotations: {}/{}\n",
            self.studio_annotation_count, self.coco_annotation_count
        ));
        s.push_str(&format!(
            "Bbox: {:.1}% matched, {:.4} avg IoU\n",
            self.bbox_validation.match_rate() * 100.0,
            self.bbox_validation.avg_iou()
        ));
        s.push_str(&format!(
            "Masks: {:.1}% preserved, {:.4} avg bbox IoU\n",
            self.mask_validation.preservation_rate() * 100.0,
            self.mask_validation.avg_bbox_iou()
        ));
        s
    }
}

impl fmt::Display for VerificationResult {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        writeln!(
            f,
            "╔══════════════════════════════════════════════════════════════╗"
        )?;
        writeln!(
            f,
            "║                  COCO IMPORT VERIFICATION                    ║"
        )?;
        writeln!(
            f,
            "╠══════════════════════════════════════════════════════════════╣"
        )?;
        writeln!(
            f,
            "║ Images:      {} in COCO, {} in Studio",
            self.coco_image_count, self.studio_image_count
        )?;
        if !self.missing_images.is_empty() {
            writeln!(f, "║   Missing:   {} images", self.missing_images.len())?;
            for name in self.missing_images.iter().take(5) {
                writeln!(f, "║              - {}", name)?;
            }
            if self.missing_images.len() > 5 {
                writeln!(
                    f,
                    "║              ... and {} more",
                    self.missing_images.len() - 5
                )?;
            }
        }
        if !self.extra_images.is_empty() {
            writeln!(f, "║   Extra:     {} images", self.extra_images.len())?;
            for name in self.extra_images.iter().take(5) {
                writeln!(f, "║              - {}", name)?;
            }
            if self.extra_images.len() > 5 {
                writeln!(
                    f,
                    "║              ... and {} more",
                    self.extra_images.len() - 5
                )?;
            }
        }
        writeln!(
            f,
            "║ Annotations: {} in COCO, {} in Studio",
            self.coco_annotation_count, self.studio_annotation_count
        )?;
        writeln!(
            f,
            "╠══════════════════════════════════════════════════════════════╣"
        )?;
        write!(f, "{}", self.bbox_validation)?;
        writeln!(
            f,
            "╠══════════════════════════════════════════════════════════════╣"
        )?;
        write!(f, "{}", self.mask_validation)?;
        writeln!(
            f,
            "╠══════════════════════════════════════════════════════════════╣"
        )?;
        write!(f, "{}", self.category_validation)?;
        writeln!(
            f,
            "╠══════════════════════════════════════════════════════════════╣"
        )?;
        let status = if self.is_valid() {
            "✓ PASSED"
        } else {
            "✗ FAILED"
        };
        writeln!(f, "║ Status: {}", status)?;
        writeln!(
            f,
            "╚══════════════════════════════════════════════════════════════╝"
        )?;
        Ok(())
    }
}

/// Bounding box validation results.
#[derive(Debug, Clone, Default)]
pub struct BboxValidationResult {
    /// Total annotations that were matched using Hungarian algorithm.
    pub total_matched: usize,
    /// Total annotations that could not be matched (IoU too low).
    pub total_unmatched: usize,
    /// Coordinate errors by range: [<1px, <2px, <5px, <10px, >=10px]
    pub errors_by_range: [usize; 5],
    /// Maximum coordinate error in pixels.
    pub max_error: f64,
    /// Sum of IoU values for averaging.
    pub sum_iou: f64,
}

impl BboxValidationResult {
    /// Returns the percentage of coordinates within 1 pixel error.
    pub fn within_1px_rate(&self) -> f64 {
        let total_coords = self.total_matched * 4;
        if total_coords == 0 {
            1.0
        } else {
            self.errors_by_range[0] as f64 / total_coords as f64
        }
    }

    /// Returns the percentage of coordinates within 2 pixels error.
    pub fn within_2px_rate(&self) -> f64 {
        let total_coords = self.total_matched * 4;
        if total_coords == 0 {
            1.0
        } else {
            (self.errors_by_range[0] + self.errors_by_range[1]) as f64 / total_coords as f64
        }
    }

    /// Returns the average IoU across all matched annotations.
    pub fn avg_iou(&self) -> f64 {
        if self.total_matched == 0 {
            1.0
        } else {
            self.sum_iou / self.total_matched as f64
        }
    }

    /// Returns the match rate (matched / total).
    pub fn match_rate(&self) -> f64 {
        let total = self.total_matched + self.total_unmatched;
        if total == 0 {
            1.0
        } else {
            self.total_matched as f64 / total as f64
        }
    }

    /// Returns true if bbox validation passes quality thresholds.
    pub fn is_valid(&self) -> bool {
        self.within_1px_rate() > 0.99 && self.match_rate() > 0.95 && self.avg_iou() > 0.95
    }
}

impl fmt::Display for BboxValidationResult {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        writeln!(f, "║ Bounding Box Validation:")?;
        writeln!(
            f,
            "║   Matched:    {}/{} ({:.1}%)",
            self.total_matched,
            self.total_matched + self.total_unmatched,
            self.match_rate() * 100.0
        )?;
        writeln!(f, "║   Avg IoU:    {:.4}", self.avg_iou())?;
        writeln!(f, "║   Within 1px: {:.1}%", self.within_1px_rate() * 100.0)?;
        writeln!(f, "║   Within 2px: {:.1}%", self.within_2px_rate() * 100.0)?;
        writeln!(f, "║   Max error:  {:.2}px", self.max_error)?;
        Ok(())
    }
}

/// Segmentation mask validation results.
#[derive(Debug, Clone, Default)]
pub struct MaskValidationResult {
    /// Annotations with segmentation in original.
    pub original_with_seg: usize,
    /// Annotations with segmentation in restored.
    pub restored_with_seg: usize,
    /// Matched pairs where both have segmentation.
    pub matched_pairs_with_seg: usize,
    /// Polygon pairs (for vertex comparison).
    pub polygon_pairs: usize,
    /// RLE pairs converted to polygon.
    pub rle_pairs: usize,
    /// Pairs where vertex count matches exactly.
    pub vertex_count_exact_match: usize,
    /// Pairs where vertex count is within 10%.
    pub vertex_count_close_match: usize,
    /// Pairs where part count matches.
    pub part_count_match: usize,
    /// Pairs with area within 1%.
    pub area_within_1pct: usize,
    /// Pairs with area within 5%.
    pub area_within_5pct: usize,
    /// Pairs with bbox IoU >= 0.9.
    pub bbox_iou_high: usize,
    /// Pairs with bbox IoU < 0.5.
    pub bbox_iou_low: usize,
    /// Sum of area ratios.
    pub sum_area_ratio: f64,
    /// Minimum area ratio.
    pub min_area_ratio: f64,
    /// Maximum area ratio.
    pub max_area_ratio: f64,
    /// Sum of bbox IoU values.
    pub sum_bbox_iou: f64,
    /// Count of zero-area segmentations.
    pub zero_area_count: usize,
}

impl MaskValidationResult {
    /// Create a new result with initialized min/max values.
    pub fn new() -> Self {
        Self {
            min_area_ratio: f64::MAX,
            max_area_ratio: 0.0,
            ..Default::default()
        }
    }

    /// Returns the segmentation preservation rate.
    pub fn preservation_rate(&self) -> f64 {
        if self.original_with_seg == 0 {
            1.0
        } else {
            self.restored_with_seg as f64 / self.original_with_seg as f64
        }
    }

    /// Returns the average area ratio.
    pub fn avg_area_ratio(&self) -> f64 {
        let valid_count = self
            .matched_pairs_with_seg
            .saturating_sub(self.zero_area_count);
        if valid_count == 0 {
            1.0
        } else {
            self.sum_area_ratio / valid_count as f64
        }
    }

    /// Returns the average bbox IoU.
    pub fn avg_bbox_iou(&self) -> f64 {
        if self.matched_pairs_with_seg == 0 {
            1.0
        } else {
            self.sum_bbox_iou / self.matched_pairs_with_seg as f64
        }
    }

    /// Returns true if mask validation passes quality thresholds.
    pub fn is_valid(&self) -> bool {
        self.preservation_rate() > 0.95 && self.avg_bbox_iou() > 0.90
    }

    /// Aggregate a single segmentation comparison into the result.
    pub fn aggregate_comparison(&mut self, cmp: &SegmentationPairComparison) {
        self.matched_pairs_with_seg += 1;

        if cmp.is_rle {
            self.rle_pairs += 1;
        } else {
            self.polygon_pairs += 1;
            if cmp.vertex_exact_match {
                self.vertex_count_exact_match += 1;
            }
            if cmp.vertex_close_match {
                self.vertex_count_close_match += 1;
            }
            if cmp.part_match {
                self.part_count_match += 1;
            }
        }

        if let Some(area_ratio) = cmp.area_ratio {
            self.sum_area_ratio += area_ratio;
            self.min_area_ratio = self.min_area_ratio.min(area_ratio);
            self.max_area_ratio = self.max_area_ratio.max(area_ratio);
            if (area_ratio - 1.0).abs() <= 0.01 {
                self.area_within_1pct += 1;
            }
            if (area_ratio - 1.0).abs() <= 0.05 {
                self.area_within_5pct += 1;
            }
        } else {
            self.zero_area_count += 1;
        }

        self.sum_bbox_iou += cmp.bbox_iou;
        if cmp.bbox_iou >= 0.9 {
            self.bbox_iou_high += 1;
        }
        if cmp.bbox_iou < 0.5 {
            self.bbox_iou_low += 1;
        }
    }
}

impl fmt::Display for MaskValidationResult {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        writeln!(f, "║ Segmentation Mask Validation:")?;
        writeln!(
            f,
            "║   Preserved:   {}/{} ({:.1}%)",
            self.restored_with_seg,
            self.original_with_seg,
            self.preservation_rate() * 100.0
        )?;
        writeln!(
            f,
            "║   Matched:     {} ({} polygon, {} RLE→polygon)",
            self.matched_pairs_with_seg, self.polygon_pairs, self.rle_pairs
        )?;
        writeln!(f, "║   Avg bbox IoU: {:.4}", self.avg_bbox_iou())?;
        writeln!(
            f,
            "║   High IoU (>=0.9): {}/{} ({:.1}%)",
            self.bbox_iou_high,
            self.matched_pairs_with_seg,
            if self.matched_pairs_with_seg > 0 {
                self.bbox_iou_high as f64 / self.matched_pairs_with_seg as f64 * 100.0
            } else {
                100.0
            }
        )?;
        if self.polygon_pairs > 0 {
            writeln!(
                f,
                "║   Vertex exact: {}/{} ({:.1}%)",
                self.vertex_count_exact_match,
                self.polygon_pairs,
                self.vertex_count_exact_match as f64 / self.polygon_pairs as f64 * 100.0
            )?;
        }
        Ok(())
    }
}

/// Category validation results.
#[derive(Debug, Clone, Default)]
pub struct CategoryValidationResult {
    /// Categories in COCO dataset.
    pub coco_categories: HashSet<String>,
    /// Categories in Studio.
    pub studio_categories: HashSet<String>,
    /// Categories missing from Studio.
    pub missing_categories: Vec<String>,
    /// Extra categories in Studio.
    pub extra_categories: Vec<String>,
}

impl CategoryValidationResult {
    /// Returns true if all categories are present.
    pub fn is_valid(&self) -> bool {
        self.missing_categories.is_empty()
    }
}

impl fmt::Display for CategoryValidationResult {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        writeln!(
            f,
            "║ Categories:    {} in COCO, {} in Studio",
            self.coco_categories.len(),
            self.studio_categories.len()
        )?;
        if !self.missing_categories.is_empty() {
            writeln!(f, "║   Missing:     {:?}", self.missing_categories)?;
        }
        if !self.extra_categories.is_empty() {
            writeln!(f, "║   Extra:       {:?}", self.extra_categories)?;
        }
        Ok(())
    }
}

/// Result of comparing a single segmentation pair.
#[derive(Debug, Clone, Default)]
pub struct SegmentationPairComparison {
    /// Whether the original was RLE (vs polygon).
    pub is_rle: bool,
    /// Vertex count exact match (polygon only).
    pub vertex_exact_match: bool,
    /// Vertex count within 10% (polygon only).
    pub vertex_close_match: bool,
    /// Part count match (polygon only).
    pub part_match: bool,
    /// Area ratio (restored / original), or None if zero area.
    pub area_ratio: Option<f64>,
    /// Bounding box IoU.
    pub bbox_iou: f64,
}

/// Compare two segmentations and return comparison metrics.
pub fn compare_segmentation_pair(
    orig_seg: &CocoSegmentation,
    rest_seg: &CocoSegmentation,
) -> SegmentationPairComparison {
    let is_rle = matches!(
        orig_seg,
        CocoSegmentation::Rle(_) | CocoSegmentation::CompressedRle(_)
    );

    let (vertex_exact_match, vertex_close_match, part_match) = if is_rle {
        (false, false, false)
    } else {
        let orig_vertices = count_polygon_vertices(orig_seg);
        let rest_vertices = count_polygon_vertices(rest_seg);
        let orig_parts = count_polygon_parts(orig_seg);
        let rest_parts = count_polygon_parts(rest_seg);

        let vertex_diff = (orig_vertices as f64 - rest_vertices as f64).abs();
        let vertex_threshold = (orig_vertices as f64 * 0.1).max(1.0);

        (
            orig_vertices == rest_vertices,
            vertex_diff <= vertex_threshold,
            orig_parts == rest_parts,
        )
    };

    let orig_area = compute_segmentation_area(orig_seg);
    let rest_area = compute_segmentation_area(rest_seg);
    let area_ratio = if orig_area > 0.0 && rest_area > 0.0 {
        Some(rest_area / orig_area)
    } else {
        None
    };

    let bbox_iou = segmentation_bbox_iou(orig_seg, rest_seg);

    SegmentationPairComparison {
        is_rle,
        vertex_exact_match,
        vertex_close_match,
        part_match,
        area_ratio,
        bbox_iou,
    }
}

/// Calculate Intersection over Union (IoU) for two COCO bboxes.
/// COCO bbox format: [x, y, width, height] (top-left corner)
pub fn bbox_iou(a: &[f64; 4], b: &[f64; 4]) -> f64 {
    let a_x1 = a[0];
    let a_y1 = a[1];
    let a_x2 = a[0] + a[2];
    let a_y2 = a[1] + a[3];

    let b_x1 = b[0];
    let b_y1 = b[1];
    let b_x2 = b[0] + b[2];
    let b_y2 = b[1] + b[3];

    // Intersection
    let inter_x1 = a_x1.max(b_x1);
    let inter_y1 = a_y1.max(b_y1);
    let inter_x2 = a_x2.min(b_x2);
    let inter_y2 = a_y2.min(b_y2);

    let inter_w = (inter_x2 - inter_x1).max(0.0);
    let inter_h = (inter_y2 - inter_y1).max(0.0);
    let inter_area = inter_w * inter_h;

    // Union
    let a_area = a[2] * a[3];
    let b_area = b[2] * b[3];
    let union_area = a_area + b_area - inter_area;

    if union_area > 0.0 {
        inter_area / union_area
    } else {
        0.0
    }
}

/// Use Hungarian algorithm to find optimal matching between two sets of
/// annotations based on bounding box IoU.
/// Returns pairs of (original_idx, restored_idx) for matched annotations.
pub fn hungarian_match<'a>(
    orig_anns: &[&'a CocoAnnotation],
    rest_anns: &[&'a CocoAnnotation],
) -> Vec<(usize, usize)> {
    if orig_anns.is_empty() || rest_anns.is_empty() {
        return vec![];
    }

    let n = orig_anns.len();
    let m = rest_anns.len();

    // Make the matrix square by padding with high-cost dummy entries
    let size = n.max(m);

    // Build cost matrix: cost = (1 - IoU) * scale
    // We use i64 for kuhn_munkres, scale by 10000 for precision
    let scale = 10000i64;
    let max_cost = scale; // Cost for non-matching (IoU = 0)

    let mut weights = Vec::with_capacity(size * size);
    for i in 0..size {
        for j in 0..size {
            let cost = match (orig_anns.get(i), rest_anns.get(j)) {
                (Some(orig), Some(rest)) => {
                    let iou = bbox_iou(&orig.bbox, &rest.bbox);
                    ((1.0 - iou) * scale as f64) as i64
                }
                _ => max_cost, // Dummy entry for padding
            };
            weights.push(cost);
        }
    }

    let matrix = Matrix::from_vec(size, size, weights).expect("Failed to create matrix");
    let (_, assignments) = kuhn_munkres_min(&matrix);

    // Filter to only real matches (not dummy) with reasonable IoU
    let min_iou_threshold = 0.3; // Only accept matches with IoU > 0.3
    assignments
        .iter()
        .enumerate()
        .filter_map(|(i, &j)| {
            if i < n && j < m {
                let iou = bbox_iou(&orig_anns[i].bbox, &rest_anns[j].bbox);
                if iou >= min_iou_threshold {
                    Some((i, j))
                } else {
                    None
                }
            } else {
                None
            }
        })
        .collect()
}

/// Calculate polygon area using the Shoelace formula.
/// Takes coordinates as flat array [x1, y1, x2, y2, ...]
pub fn polygon_area(coords: &[f64]) -> f64 {
    let n = coords.len() / 2;
    if n < 3 {
        return 0.0;
    }

    let mut area = 0.0;
    for i in 0..n {
        let j = (i + 1) % n;
        let x_i = coords[i * 2];
        let y_i = coords[i * 2 + 1];
        let x_j = coords[j * 2];
        let y_j = coords[j * 2 + 1];
        area += x_i * y_j - x_j * y_i;
    }
    (area / 2.0).abs()
}

/// Calculate total area of a segmentation.
pub fn compute_segmentation_area(seg: &CocoSegmentation) -> f64 {
    match seg {
        CocoSegmentation::Polygon(polys) => polys.iter().map(|p| polygon_area(p)).sum(),
        CocoSegmentation::Rle(rle) => {
            if let Ok((mask, _, _)) = decode_rle(rle) {
                mask.iter().filter(|&&v| v == 1).count() as f64
            } else {
                0.0
            }
        }
        CocoSegmentation::CompressedRle(compressed) => {
            if let Ok((mask, _, _)) = decode_compressed_rle(compressed) {
                mask.iter().filter(|&&v| v == 1).count() as f64
            } else {
                0.0
            }
        }
    }
}

/// Calculate bounding box of a polygon (min_x, min_y, max_x, max_y)
pub fn polygon_bounds(coords: &[f64]) -> Option<(f64, f64, f64, f64)> {
    if coords.len() < 4 {
        return None;
    }
    let mut min_x = f64::MAX;
    let mut min_y = f64::MAX;
    let mut max_x = f64::MIN;
    let mut max_y = f64::MIN;

    for chunk in coords.chunks(2) {
        if chunk.len() == 2 {
            min_x = min_x.min(chunk[0]);
            max_x = max_x.max(chunk[0]);
            min_y = min_y.min(chunk[1]);
            max_y = max_y.max(chunk[1]);
        }
    }
    Some((min_x, min_y, max_x, max_y))
}

/// Get bounding box for any segmentation type.
pub fn segmentation_bounds(seg: &CocoSegmentation) -> Option<(f64, f64, f64, f64)> {
    match seg {
        CocoSegmentation::Polygon(polys) => {
            polys
                .iter()
                .filter_map(|p| polygon_bounds(p))
                .fold(None, |acc, b| match acc {
                    None => Some(b),
                    Some((min_x, min_y, max_x, max_y)) => Some((
                        min_x.min(b.0),
                        min_y.min(b.1),
                        max_x.max(b.2),
                        max_y.max(b.3),
                    )),
                })
        }
        CocoSegmentation::Rle(rle) => {
            let (mask, height, width) = decode_rle(rle).ok()?;
            rle_mask_bounds(&mask, height, width)
        }
        CocoSegmentation::CompressedRle(compressed) => {
            let (mask, height, width) = decode_compressed_rle(compressed).ok()?;
            rle_mask_bounds(&mask, height, width)
        }
    }
}

/// Find bounds of a binary mask.
fn rle_mask_bounds(mask: &[u8], height: u32, width: u32) -> Option<(f64, f64, f64, f64)> {
    let mut min_x = width;
    let mut min_y = height;
    let mut max_x = 0u32;
    let mut max_y = 0u32;
    let mut found_any = false;

    for y in 0..height {
        for x in 0..width {
            let idx = (y as usize) * (width as usize) + (x as usize);
            if mask.get(idx) == Some(&1) {
                found_any = true;
                min_x = min_x.min(x);
                max_x = max_x.max(x);
                min_y = min_y.min(y);
                max_y = max_y.max(y);
            }
        }
    }

    if found_any {
        Some((min_x as f64, min_y as f64, max_x as f64, max_y as f64))
    } else {
        None
    }
}

/// Calculate IoU between two segmentation bounding boxes.
pub fn segmentation_bbox_iou(seg1: &CocoSegmentation, seg2: &CocoSegmentation) -> f64 {
    let bounds1 = segmentation_bounds(seg1);
    let bounds2 = segmentation_bounds(seg2);

    match (bounds1, bounds2) {
        (Some((a_x1, a_y1, a_x2, a_y2)), Some((b_x1, b_y1, b_x2, b_y2))) => {
            let inter_x1 = a_x1.max(b_x1);
            let inter_y1 = a_y1.max(b_y1);
            let inter_x2 = a_x2.min(b_x2);
            let inter_y2 = a_y2.min(b_y2);

            let inter_w = (inter_x2 - inter_x1).max(0.0);
            let inter_h = (inter_y2 - inter_y1).max(0.0);
            let inter_area = inter_w * inter_h;

            let a_area = (a_x2 - a_x1) * (a_y2 - a_y1);
            let b_area = (b_x2 - b_x1) * (b_y2 - b_y1);
            let union_area = a_area + b_area - inter_area;

            if union_area > 0.0 {
                inter_area / union_area
            } else {
                0.0
            }
        }
        _ => 0.0,
    }
}

/// Count total polygon vertices in a segmentation.
pub fn count_polygon_vertices(seg: &CocoSegmentation) -> usize {
    match seg {
        CocoSegmentation::Polygon(polys) => polys.iter().map(|p| p.len() / 2).sum(),
        _ => 0,
    }
}

/// Count number of polygon parts in a segmentation.
pub fn count_polygon_parts(seg: &CocoSegmentation) -> usize {
    match seg {
        CocoSegmentation::Polygon(polys) => polys.len(),
        _ => 0,
    }
}

/// Build a map of annotations by sample name for efficient lookup.
pub fn build_annotation_map_by_name(
    dataset: &CocoDataset,
) -> HashMap<String, Vec<&CocoAnnotation>> {
    let image_names: HashMap<u64, String> = dataset
        .images
        .iter()
        .map(|img| {
            let name = std::path::Path::new(&img.file_name)
                .file_stem()
                .and_then(|s| s.to_str())
                .unwrap_or(&img.file_name)
                .to_string();
            (img.id, name)
        })
        .collect();

    let mut map: HashMap<String, Vec<_>> = HashMap::new();
    for ann in &dataset.annotations {
        if let Some(name) = image_names.get(&ann.image_id) {
            map.entry(name.clone()).or_default().push(ann);
        }
    }
    map
}

/// Validate bounding boxes between two datasets using Hungarian matching.
pub fn validate_bboxes(original: &CocoDataset, restored: &CocoDataset) -> BboxValidationResult {
    let mut result = BboxValidationResult::default();

    let original_by_name = build_annotation_map_by_name(original);
    let restored_by_name = build_annotation_map_by_name(restored);

    for (name, orig_anns) in &original_by_name {
        if let Some(rest_anns) = restored_by_name.get(name) {
            let matches = hungarian_match(orig_anns, rest_anns);

            for (orig_idx, rest_idx) in &matches {
                let orig_ann = orig_anns[*orig_idx];
                let rest_ann = rest_anns[*rest_idx];

                // Track IoU
                let iou = bbox_iou(&orig_ann.bbox, &rest_ann.bbox);
                result.sum_iou += iou;

                // Measure coordinate errors
                for i in 0..4 {
                    let error = (orig_ann.bbox[i] - rest_ann.bbox[i]).abs();
                    result.max_error = result.max_error.max(error);

                    if error < 1.0 {
                        result.errors_by_range[0] += 1;
                    } else if error < 2.0 {
                        result.errors_by_range[1] += 1;
                    } else if error < 5.0 {
                        result.errors_by_range[2] += 1;
                    } else if error < 10.0 {
                        result.errors_by_range[3] += 1;
                    } else {
                        result.errors_by_range[4] += 1;
                    }
                }
                result.total_matched += 1;
            }

            result.total_unmatched += orig_anns.len() - matches.len();
        } else {
            result.total_unmatched += orig_anns.len();
        }
    }

    result
}

/// Validate segmentation masks between two datasets using Hungarian matching.
pub fn validate_masks(original: &CocoDataset, restored: &CocoDataset) -> MaskValidationResult {
    let mut result = MaskValidationResult::new();

    // Count segmentations in original and restored
    result.original_with_seg = original
        .annotations
        .iter()
        .filter(|a| a.segmentation.is_some())
        .count();
    result.restored_with_seg = restored
        .annotations
        .iter()
        .filter(|a| a.segmentation.is_some())
        .count();

    let original_by_name = build_annotation_map_by_name(original);
    let restored_by_name = build_annotation_map_by_name(restored);

    for (name, orig_anns) in &original_by_name {
        if let Some(rest_anns) = restored_by_name.get(name) {
            let matches = hungarian_match(orig_anns, rest_anns);

            for (orig_idx, rest_idx) in &matches {
                let orig_ann = orig_anns[*orig_idx];
                let rest_ann = rest_anns[*rest_idx];

                if let (Some(orig_seg), Some(rest_seg)) =
                    (&orig_ann.segmentation, &rest_ann.segmentation)
                {
                    let comparison = compare_segmentation_pair(orig_seg, rest_seg);
                    result.aggregate_comparison(&comparison);
                }
            }
        }
    }

    result
}

/// Validate categories between two datasets.
pub fn validate_categories(
    original: &CocoDataset,
    restored: &CocoDataset,
) -> CategoryValidationResult {
    let coco_cats: HashSet<String> = original.categories.iter().map(|c| c.name.clone()).collect();
    let studio_cats: HashSet<String> = restored.categories.iter().map(|c| c.name.clone()).collect();

    let missing: Vec<String> = coco_cats.difference(&studio_cats).cloned().collect();
    let extra: Vec<String> = studio_cats.difference(&coco_cats).cloned().collect();

    CategoryValidationResult {
        coco_categories: coco_cats,
        studio_categories: studio_cats,
        missing_categories: missing,
        extra_categories: extra,
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::coco::{CocoCategory, CocoImage, CocoRle};

    // =========================================================================
    // bbox_iou tests
    // =========================================================================

    #[test]
    fn test_bbox_iou_perfect_overlap() {
        let a = [0.0, 0.0, 100.0, 100.0];
        let b = [0.0, 0.0, 100.0, 100.0];
        assert!((bbox_iou(&a, &b) - 1.0).abs() < 1e-6);
    }

    #[test]
    fn test_bbox_iou_no_overlap() {
        let a = [0.0, 0.0, 100.0, 100.0];
        let b = [200.0, 200.0, 100.0, 100.0];
        assert!(bbox_iou(&a, &b) < 1e-6);
    }

    #[test]
    fn test_bbox_iou_partial_overlap() {
        let a = [0.0, 0.0, 100.0, 100.0];
        let b = [50.0, 50.0, 100.0, 100.0];
        // Intersection: 50x50 = 2500, Union: 10000 + 10000 - 2500 = 17500
        let expected = 2500.0 / 17500.0;
        assert!((bbox_iou(&a, &b) - expected).abs() < 1e-6);
    }

    #[test]
    fn test_bbox_iou_contained() {
        // b is fully contained in a
        let a = [0.0, 0.0, 100.0, 100.0];
        let b = [25.0, 25.0, 50.0, 50.0];
        // Intersection: 50x50 = 2500, Union: 10000
        let expected = 2500.0 / 10000.0;
        assert!((bbox_iou(&a, &b) - expected).abs() < 1e-6);
    }

    #[test]
    fn test_bbox_iou_zero_area() {
        let a = [0.0, 0.0, 0.0, 0.0];
        let b = [0.0, 0.0, 100.0, 100.0];
        assert!(bbox_iou(&a, &b) < 1e-6);
    }

    // =========================================================================
    // polygon_area tests
    // =========================================================================

    #[test]
    fn test_polygon_area_square() {
        // 10x10 square
        let coords = [0.0, 0.0, 10.0, 0.0, 10.0, 10.0, 0.0, 10.0];
        assert!((polygon_area(&coords) - 100.0).abs() < 1e-6);
    }

    #[test]
    fn test_polygon_area_triangle() {
        // Triangle with base 10 and height 10
        let coords = [0.0, 0.0, 10.0, 0.0, 5.0, 10.0];
        assert!((polygon_area(&coords) - 50.0).abs() < 1e-6);
    }

    #[test]
    fn test_polygon_area_too_small() {
        // Less than 3 points (6 coords)
        let coords = [0.0, 0.0, 10.0, 10.0];
        assert!(polygon_area(&coords) < 1e-6);
    }

    #[test]
    fn test_polygon_area_complex() {
        // L-shaped polygon (can compute as two rectangles)
        // Points: (0,0), (20,0), (20,10), (10,10), (10,20), (0,20)
        let coords = [
            0.0, 0.0, 20.0, 0.0, 20.0, 10.0, 10.0, 10.0, 10.0, 20.0, 0.0, 20.0,
        ];
        // Area = 10*20 + 10*10 = 300
        assert!((polygon_area(&coords) - 300.0).abs() < 1e-6);
    }

    // =========================================================================
    // polygon_bounds tests
    // =========================================================================

    #[test]
    fn test_polygon_bounds_square() {
        let coords = [0.0, 0.0, 100.0, 0.0, 100.0, 100.0, 0.0, 100.0];
        let bounds = polygon_bounds(&coords);
        assert_eq!(bounds, Some((0.0, 0.0, 100.0, 100.0)));
    }

    #[test]
    fn test_polygon_bounds_offset() {
        let coords = [50.0, 60.0, 150.0, 60.0, 150.0, 160.0, 50.0, 160.0];
        let bounds = polygon_bounds(&coords);
        assert_eq!(bounds, Some((50.0, 60.0, 150.0, 160.0)));
    }

    #[test]
    fn test_polygon_bounds_too_small() {
        let coords = [0.0, 0.0];
        assert!(polygon_bounds(&coords).is_none());
    }

    // =========================================================================
    // hungarian_match tests
    // =========================================================================

    #[test]
    fn test_hungarian_match_empty_inputs() {
        let orig: Vec<&CocoAnnotation> = vec![];
        let rest: Vec<&CocoAnnotation> = vec![];
        let matches = hungarian_match(&orig, &rest);
        assert!(matches.is_empty());
    }

    #[test]
    fn test_hungarian_match_perfect_match() {
        let ann1 = CocoAnnotation {
            id: 1,
            image_id: 1,
            category_id: 1,
            bbox: [0.0, 0.0, 100.0, 100.0],
            ..Default::default()
        };
        let ann2 = CocoAnnotation {
            id: 2,
            image_id: 1,
            category_id: 1,
            bbox: [0.0, 0.0, 100.0, 100.0], // Same bbox
            ..Default::default()
        };

        let orig = vec![&ann1];
        let rest = vec![&ann2];
        let matches = hungarian_match(&orig, &rest);

        assert_eq!(matches.len(), 1);
        assert_eq!(matches[0], (0, 0));
    }

    #[test]
    fn test_hungarian_match_multiple() {
        let ann1 = CocoAnnotation {
            id: 1,
            image_id: 1,
            category_id: 1,
            bbox: [0.0, 0.0, 50.0, 50.0],
            ..Default::default()
        };
        let ann2 = CocoAnnotation {
            id: 2,
            image_id: 1,
            category_id: 1,
            bbox: [100.0, 100.0, 50.0, 50.0],
            ..Default::default()
        };

        let ann3 = CocoAnnotation {
            id: 3,
            image_id: 1,
            category_id: 1,
            bbox: [100.0, 100.0, 50.0, 50.0], // Matches ann2
            ..Default::default()
        };
        let ann4 = CocoAnnotation {
            id: 4,
            image_id: 1,
            category_id: 1,
            bbox: [0.0, 0.0, 50.0, 50.0], // Matches ann1
            ..Default::default()
        };

        let orig = vec![&ann1, &ann2];
        let rest = vec![&ann3, &ann4];
        let matches = hungarian_match(&orig, &rest);

        assert_eq!(matches.len(), 2);
    }

    #[test]
    fn test_hungarian_match_unequal_sizes() {
        let ann1 = CocoAnnotation {
            id: 1,
            image_id: 1,
            category_id: 1,
            bbox: [0.0, 0.0, 100.0, 100.0],
            ..Default::default()
        };
        let ann2 = CocoAnnotation {
            id: 2,
            image_id: 1,
            category_id: 1,
            bbox: [200.0, 200.0, 100.0, 100.0],
            ..Default::default()
        };
        let ann3 = CocoAnnotation {
            id: 3,
            image_id: 1,
            category_id: 1,
            bbox: [0.0, 0.0, 100.0, 100.0], // Matches ann1
            ..Default::default()
        };

        let orig = vec![&ann1, &ann2];
        let rest = vec![&ann3]; // Fewer restored
        let matches = hungarian_match(&orig, &rest);

        assert_eq!(matches.len(), 1);
        assert_eq!(matches[0], (0, 0)); // ann1 matched to ann3
    }

    // =========================================================================
    // BboxValidationResult tests
    // =========================================================================

    #[test]
    fn test_bbox_validation_result_rates() {
        let mut result = BboxValidationResult::default();
        result.total_matched = 100;
        result.total_unmatched = 10;
        result.errors_by_range[0] = 350; // 350/400 = 87.5%
        result.errors_by_range[1] = 40;
        result.sum_iou = 95.0;

        assert!((result.match_rate() - 0.909).abs() < 0.01);
        assert!((result.avg_iou() - 0.95).abs() < 0.01);
    }

    #[test]
    fn test_bbox_validation_result_empty() {
        let result = BboxValidationResult::default();
        assert!((result.match_rate() - 1.0).abs() < 1e-6);
        assert!((result.avg_iou() - 1.0).abs() < 1e-6);
        assert!((result.within_1px_rate() - 1.0).abs() < 1e-6);
    }

    #[test]
    fn test_bbox_validation_result_is_valid() {
        let mut result = BboxValidationResult::default();
        result.total_matched = 100;
        result.errors_by_range[0] = 400; // All within 1px
        result.sum_iou = 98.0;
        assert!(result.is_valid());
    }

    #[test]
    fn test_bbox_validation_result_not_valid() {
        let mut result = BboxValidationResult::default();
        result.total_matched = 100;
        result.total_unmatched = 50; // Low match rate
        result.errors_by_range[0] = 200;
        result.sum_iou = 50.0;
        assert!(!result.is_valid());
    }

    // =========================================================================
    // MaskValidationResult tests
    // =========================================================================

    #[test]
    fn test_mask_validation_result_new() {
        let result = MaskValidationResult::new();
        assert_eq!(result.min_area_ratio, f64::MAX);
        assert_eq!(result.max_area_ratio, 0.0);
    }

    #[test]
    fn test_mask_validation_result_preservation_rate() {
        let mut result = MaskValidationResult::new();
        result.original_with_seg = 100;
        result.restored_with_seg = 95;
        assert!((result.preservation_rate() - 0.95).abs() < 1e-6);
    }

    #[test]
    fn test_mask_validation_result_empty() {
        let result = MaskValidationResult::new();
        assert!((result.preservation_rate() - 1.0).abs() < 1e-6);
        assert!((result.avg_area_ratio() - 1.0).abs() < 1e-6);
        assert!((result.avg_bbox_iou() - 1.0).abs() < 1e-6);
    }

    // =========================================================================
    // CategoryValidationResult tests
    // =========================================================================

    #[test]
    fn test_category_validation_result_is_valid() {
        let result = CategoryValidationResult {
            coco_categories: ["person", "car"].iter().map(|s| s.to_string()).collect(),
            studio_categories: ["person", "car"].iter().map(|s| s.to_string()).collect(),
            missing_categories: vec![],
            extra_categories: vec![],
        };
        assert!(result.is_valid());
    }

    #[test]
    fn test_category_validation_result_missing() {
        let result = CategoryValidationResult {
            coco_categories: ["person", "car"].iter().map(|s| s.to_string()).collect(),
            studio_categories: ["person"].iter().map(|s| s.to_string()).collect(),
            missing_categories: vec!["car".to_string()],
            extra_categories: vec![],
        };
        assert!(!result.is_valid());
    }

    // =========================================================================
    // compute_segmentation_area tests
    // =========================================================================

    #[test]
    fn test_compute_segmentation_area_polygon() {
        let seg =
            CocoSegmentation::Polygon(vec![vec![0.0, 0.0, 100.0, 0.0, 100.0, 100.0, 0.0, 100.0]]);
        let area = compute_segmentation_area(&seg);
        assert!((area - 10000.0).abs() < 1e-6);
    }

    #[test]
    fn test_compute_segmentation_area_multiple_polygons() {
        // Two 10x10 squares
        let seg = CocoSegmentation::Polygon(vec![
            vec![0.0, 0.0, 10.0, 0.0, 10.0, 10.0, 0.0, 10.0],
            vec![20.0, 20.0, 30.0, 20.0, 30.0, 30.0, 20.0, 30.0],
        ]);
        let area = compute_segmentation_area(&seg);
        assert!((area - 200.0).abs() < 1e-6);
    }

    // =========================================================================
    // count_polygon_vertices/parts tests
    // =========================================================================

    #[test]
    fn test_count_polygon_vertices() {
        let seg = CocoSegmentation::Polygon(vec![vec![0.0, 0.0, 10.0, 0.0, 10.0, 10.0, 0.0, 10.0]]);
        assert_eq!(count_polygon_vertices(&seg), 4);
    }

    #[test]
    fn test_count_polygon_vertices_multiple() {
        let seg = CocoSegmentation::Polygon(vec![
            vec![0.0, 0.0, 10.0, 0.0, 10.0, 10.0], // 3 vertices
            vec![20.0, 20.0, 30.0, 20.0, 30.0, 30.0, 20.0, 30.0], // 4 vertices
        ]);
        assert_eq!(count_polygon_vertices(&seg), 7);
    }

    #[test]
    fn test_count_polygon_parts() {
        let seg = CocoSegmentation::Polygon(vec![
            vec![0.0, 0.0, 10.0, 0.0, 10.0, 10.0],
            vec![20.0, 20.0, 30.0, 20.0, 30.0, 30.0],
        ]);
        assert_eq!(count_polygon_parts(&seg), 2);
    }

    #[test]
    fn test_count_polygon_vertices_rle() {
        let rle = CocoRle {
            counts: vec![100],
            size: [10, 10],
        };
        let seg = CocoSegmentation::Rle(rle);
        assert_eq!(count_polygon_vertices(&seg), 0);
    }

    // =========================================================================
    // VerificationResult tests
    // =========================================================================

    #[test]
    fn test_verification_result_is_valid() {
        let result = VerificationResult {
            coco_image_count: 100,
            studio_image_count: 100,
            missing_images: vec![],
            extra_images: vec![],
            coco_annotation_count: 500,
            studio_annotation_count: 500,
            bbox_validation: {
                let mut bv = BboxValidationResult::default();
                bv.total_matched = 500;
                bv.errors_by_range[0] = 2000; // All within 1px
                bv.sum_iou = 495.0;
                bv
            },
            mask_validation: {
                let mut mv = MaskValidationResult::new();
                mv.original_with_seg = 500;
                mv.restored_with_seg = 500;
                mv.matched_pairs_with_seg = 500;
                mv.sum_bbox_iou = 475.0;
                mv
            },
            category_validation: CategoryValidationResult {
                coco_categories: ["person"].iter().map(|s| s.to_string()).collect(),
                studio_categories: ["person"].iter().map(|s| s.to_string()).collect(),
                missing_categories: vec![],
                extra_categories: vec![],
            },
        };

        assert!(result.is_valid());
    }

    #[test]
    fn test_verification_result_summary() {
        let result = VerificationResult {
            coco_image_count: 100,
            studio_image_count: 98,
            missing_images: vec!["img1.jpg".to_string(), "img2.jpg".to_string()],
            extra_images: vec![],
            coco_annotation_count: 500,
            studio_annotation_count: 490,
            bbox_validation: BboxValidationResult::default(),
            mask_validation: MaskValidationResult::new(),
            category_validation: CategoryValidationResult::default(),
        };

        let summary = result.summary();
        assert!(summary.contains("Images:"));
        assert!(summary.contains("Annotations:"));
    }

    // =========================================================================
    // build_annotation_map_by_name tests
    // =========================================================================

    #[test]
    fn test_build_annotation_map_by_name() {
        let dataset = CocoDataset {
            images: vec![
                CocoImage {
                    id: 1,
                    file_name: "image1.jpg".to_string(),
                    ..Default::default()
                },
                CocoImage {
                    id: 2,
                    file_name: "image2.jpg".to_string(),
                    ..Default::default()
                },
            ],
            annotations: vec![
                CocoAnnotation {
                    id: 1,
                    image_id: 1,
                    ..Default::default()
                },
                CocoAnnotation {
                    id: 2,
                    image_id: 1,
                    ..Default::default()
                },
                CocoAnnotation {
                    id: 3,
                    image_id: 2,
                    ..Default::default()
                },
            ],
            ..Default::default()
        };

        let map = build_annotation_map_by_name(&dataset);

        assert_eq!(map.len(), 2);
        assert_eq!(map.get("image1").unwrap().len(), 2);
        assert_eq!(map.get("image2").unwrap().len(), 1);
    }

    // =========================================================================
    // validate_categories tests
    // =========================================================================

    #[test]
    fn test_validate_categories_match() {
        let original = CocoDataset {
            categories: vec![
                CocoCategory {
                    id: 1,
                    name: "cat".to_string(),
                    supercategory: None,
                },
                CocoCategory {
                    id: 2,
                    name: "dog".to_string(),
                    supercategory: None,
                },
            ],
            ..Default::default()
        };

        let restored = CocoDataset {
            categories: vec![
                CocoCategory {
                    id: 1,
                    name: "cat".to_string(),
                    supercategory: None,
                },
                CocoCategory {
                    id: 2,
                    name: "dog".to_string(),
                    supercategory: None,
                },
            ],
            ..Default::default()
        };

        let result = validate_categories(&original, &restored);
        assert!(result.is_valid());
        assert!(result.missing_categories.is_empty());
        assert!(result.extra_categories.is_empty());
    }

    #[test]
    fn test_validate_categories_missing_and_extra() {
        let original = CocoDataset {
            categories: vec![
                CocoCategory {
                    id: 1,
                    name: "cat".to_string(),
                    supercategory: None,
                },
                CocoCategory {
                    id: 2,
                    name: "dog".to_string(),
                    supercategory: None,
                },
            ],
            ..Default::default()
        };

        let restored = CocoDataset {
            categories: vec![
                CocoCategory {
                    id: 1,
                    name: "cat".to_string(),
                    supercategory: None,
                },
                CocoCategory {
                    id: 3,
                    name: "bird".to_string(),
                    supercategory: None,
                },
            ],
            ..Default::default()
        };

        let result = validate_categories(&original, &restored);
        assert!(!result.is_valid());
        assert!(result.missing_categories.contains(&"dog".to_string()));
        assert!(result.extra_categories.contains(&"bird".to_string()));
    }

    // =========================================================================
    // compare_segmentation_pair tests
    // =========================================================================

    #[test]
    fn test_compare_segmentation_pair_identical_polygons() {
        let seg =
            CocoSegmentation::Polygon(vec![vec![0.0, 0.0, 100.0, 0.0, 100.0, 100.0, 0.0, 100.0]]);
        let result = compare_segmentation_pair(&seg, &seg);

        assert!(!result.is_rle);
        assert!(result.vertex_exact_match);
        assert!(result.vertex_close_match);
        assert!(result.part_match);
        assert!(result.area_ratio.is_some());
        assert!((result.area_ratio.unwrap() - 1.0).abs() < 1e-6);
        assert!((result.bbox_iou - 1.0).abs() < 1e-6);
    }

    #[test]
    fn test_compare_segmentation_pair_different_vertex_count() {
        let seg1 =
            CocoSegmentation::Polygon(vec![vec![0.0, 0.0, 100.0, 0.0, 100.0, 100.0, 0.0, 100.0]]);
        // Triangle with same bounding box
        let seg2 = CocoSegmentation::Polygon(vec![vec![0.0, 0.0, 100.0, 0.0, 50.0, 100.0]]);

        let result = compare_segmentation_pair(&seg1, &seg2);

        assert!(!result.is_rle);
        assert!(!result.vertex_exact_match); // 4 vs 3 vertices
        // Note: vertex_close_match is true because threshold is max(10%, 1) = 1, and
        // diff = 1
        assert!(result.vertex_close_match);
        assert!(result.part_match); // Both have 1 part
    }

    #[test]
    fn test_compare_segmentation_pair_rle() {
        let rle = CocoRle {
            counts: vec![100],
            size: [10, 10],
        };
        let seg = CocoSegmentation::Rle(rle);
        let poly =
            CocoSegmentation::Polygon(vec![vec![0.0, 0.0, 10.0, 0.0, 10.0, 10.0, 0.0, 10.0]]);

        let result = compare_segmentation_pair(&seg, &poly);

        assert!(result.is_rle);
        assert!(!result.vertex_exact_match);
        assert!(!result.vertex_close_match);
        assert!(!result.part_match);
    }

    #[test]
    fn test_compare_segmentation_pair_scaled() {
        let seg1 =
            CocoSegmentation::Polygon(vec![vec![0.0, 0.0, 100.0, 0.0, 100.0, 100.0, 0.0, 100.0]]);
        let seg2 =
            CocoSegmentation::Polygon(vec![vec![0.0, 0.0, 50.0, 0.0, 50.0, 50.0, 0.0, 50.0]]);

        let result = compare_segmentation_pair(&seg1, &seg2);

        assert!(result.area_ratio.is_some());
        // 2500 / 10000 = 0.25
        assert!((result.area_ratio.unwrap() - 0.25).abs() < 0.01);
    }

    // =========================================================================
    // aggregate_comparison tests
    // =========================================================================

    #[test]
    fn test_aggregate_comparison_polygon() {
        let mut result = MaskValidationResult::new();
        let cmp = SegmentationPairComparison {
            is_rle: false,
            vertex_exact_match: true,
            vertex_close_match: true,
            part_match: true,
            area_ratio: Some(1.0),
            bbox_iou: 0.95,
        };

        result.aggregate_comparison(&cmp);

        assert_eq!(result.matched_pairs_with_seg, 1);
        assert_eq!(result.polygon_pairs, 1);
        assert_eq!(result.rle_pairs, 0);
        assert_eq!(result.vertex_count_exact_match, 1);
        assert_eq!(result.vertex_count_close_match, 1);
        assert_eq!(result.part_count_match, 1);
        assert_eq!(result.area_within_1pct, 1);
        assert_eq!(result.area_within_5pct, 1);
        assert_eq!(result.bbox_iou_high, 1);
    }

    #[test]
    fn test_aggregate_comparison_rle() {
        let mut result = MaskValidationResult::new();
        let cmp = SegmentationPairComparison {
            is_rle: true,
            vertex_exact_match: false,
            vertex_close_match: false,
            part_match: false,
            area_ratio: Some(0.98),
            bbox_iou: 0.92,
        };

        result.aggregate_comparison(&cmp);

        assert_eq!(result.matched_pairs_with_seg, 1);
        assert_eq!(result.polygon_pairs, 0);
        assert_eq!(result.rle_pairs, 1);
        assert_eq!(result.vertex_count_exact_match, 0);
        assert_eq!(result.area_within_5pct, 1);
        assert_eq!(result.bbox_iou_high, 1);
    }

    #[test]
    fn test_aggregate_comparison_zero_area() {
        let mut result = MaskValidationResult::new();
        let cmp = SegmentationPairComparison {
            is_rle: false,
            vertex_exact_match: true,
            vertex_close_match: true,
            part_match: true,
            area_ratio: None, // Zero area
            bbox_iou: 0.3,
        };

        result.aggregate_comparison(&cmp);

        assert_eq!(result.zero_area_count, 1);
        assert_eq!(result.area_within_1pct, 0);
        assert_eq!(result.bbox_iou_low, 1);
    }

    #[test]
    fn test_aggregate_comparison_multiple() {
        let mut result = MaskValidationResult::new();

        let cmp1 = SegmentationPairComparison {
            is_rle: false,
            vertex_exact_match: true,
            vertex_close_match: true,
            part_match: true,
            area_ratio: Some(1.0),
            bbox_iou: 0.95,
        };
        let cmp2 = SegmentationPairComparison {
            is_rle: true,
            vertex_exact_match: false,
            vertex_close_match: false,
            part_match: false,
            area_ratio: Some(0.9),
            bbox_iou: 0.85,
        };

        result.aggregate_comparison(&cmp1);
        result.aggregate_comparison(&cmp2);

        assert_eq!(result.matched_pairs_with_seg, 2);
        assert_eq!(result.polygon_pairs, 1);
        assert_eq!(result.rle_pairs, 1);
        assert!((result.sum_area_ratio - 1.9).abs() < 0.01);
        assert!((result.min_area_ratio - 0.9).abs() < 0.01);
        assert!((result.max_area_ratio - 1.0).abs() < 0.01);
    }
}