kreuzberg 4.5.3

//! Segment-to-region spatial assignment logic.

use crate::pdf::hierarchy::SegmentData;

use super::super::geometry::Rect;
use super::LayoutRegion;
use crate::pdf::markdown::types::{LayoutHint, LayoutHintClass};

/// Minimum intersection-over-self (IoS) for a segment to be assigned to a region.
/// Matches the reference threshold of 0.2 (20% of the segment's area must overlap).
const MIN_IOS_THRESHOLD: f32 = 0.2;

/// Minimum alphanumeric character count for Picture region text to be considered
/// substantive. Below this threshold, text is treated as diagram labels or
/// decorative and suppressed. Above it, the text is preserved as unassigned
/// (e.g., screenshots of papers, appendix pages containing readable text).
const PICTURE_SUBSTANTIVE_CHAR_THRESHOLD: usize = 50;

/// Padding (in points) added to each side of the tight bbox during refinement.
const BBOX_REFINEMENT_PADDING: f32 = 2.0;

/// Maximum number of assign→shrink→re-assign iterations.
const MAX_REFINEMENT_ITERATIONS: usize = 3;

/// Assign each segment to its best-matching layout region.
///
/// Uses intersection-over-self (IoS): the fraction of the segment's bounding box
/// that overlaps with a region. If IoS >= 0.2, the segment can be assigned.
/// Among qualifying regions, the one with highest IoS wins; ties broken by
/// smallest area (most specific region).
///
/// Table and Picture regions are excluded from assignment — handled separately.
/// Segments overlapping successfully extracted tables are suppressed (>=50% IoS).
/// Segments overlapping Picture regions are evaluated per-region: regions
/// validated as empty (no text CCs) never suppress. Regions with substantive
/// text (>= 50 alphanumeric chars) are preserved as unassigned, allowing
/// screenshots and misclassified text-heavy regions to flow through the
/// standard pipeline. Only short/label text (diagram labels, axis text) is
/// suppressed.
pub(in crate::pdf::markdown) fn assign_segments_to_regions<'a>(
    segments: &[SegmentData],
    hints: &'a [LayoutHint],
    min_confidence: f32,
    extracted_table_bboxes: &[crate::types::BoundingBox],
    hint_validations: &[super::layout_validation::RegionValidation],
) -> (Vec<LayoutRegion<'a>>, Vec<usize>) {
    let confident_hints: Vec<&LayoutHint> = hints
        .iter()
        .filter(|h| h.confidence >= min_confidence)
        // Exclude Table and Picture — Table text is handled by TATR extraction,
        // Picture text is suppressed or preserved based on substantive text check.
        .filter(|h| !matches!(h.class, LayoutHintClass::Table | LayoutHintClass::Picture))
        .collect();

    // Suppress segments overlapping successfully extracted tables (>=50% IoS).
    // Only tables that actually produced TATR output have entries in
    // extracted_table_bboxes — failed extractions don't suppress anything.
    let suppress_bboxes = extracted_table_bboxes;

    // Collect Picture region bounding boxes for suppression, paired with
    // their validation status. Empty-validated Picture regions don't suppress.
    let picture_hints: Vec<(&LayoutHint, bool)> = hints
        .iter()
        .enumerate()
        .filter(|(_, h)| h.confidence >= min_confidence && h.class == LayoutHintClass::Picture)
        .map(|(idx, h)| {
            let is_empty = hint_validations
                .get(idx)
                .is_some_and(|v| *v == super::layout_validation::RegionValidation::Empty);
            (h, is_empty)
        })
        .collect();

    if confident_hints.is_empty() && suppress_bboxes.is_empty() && picture_hints.is_empty() {
        let all_indices: Vec<usize> = (0..segments.len()).collect();
        return (Vec::new(), all_indices);
    }

    // Pre-compute hint areas for tie-breaking
    let hint_areas: Vec<f32> = confident_hints
        .iter()
        .map(|h| (h.right - h.left) * (h.top - h.bottom))
        .collect();

    // Build region containers
    let mut regions: Vec<LayoutRegion> = confident_hints
        .iter()
        .map(|&hint| LayoutRegion {
            hint,
            segment_indices: Vec::new(),
            merged_bbox: None,
        })
        .collect();

    let mut unassigned: Vec<usize> = Vec::new();

    let mut suppressed_count = 0_usize;

    // First pass: identify which segments fall into which Picture region.
    // We collect per-picture segment indices to decide per-region whether to
    // suppress (decorative/label text) or preserve (substantive text like
    // screenshots of papers).
    let mut picture_seg_indices: Vec<Vec<usize>> = vec![Vec::new(); picture_hints.len()];

    for (seg_idx, seg) in segments.iter().enumerate() {
        if seg.text.trim().is_empty() {
            continue;
        }
        let seg_rect = Rect::from_lbrt(seg.x, seg.y, seg.x + seg.width, seg.y + seg.height);
        for (pi, (ph, _)) in picture_hints.iter().enumerate() {
            let hint_rect = Rect::from_lbrt(ph.left, ph.bottom, ph.right, ph.top);
            if seg_rect.intersection_over_self(&hint_rect) >= 0.5 {
                picture_seg_indices[pi].push(seg_idx);
                break; // Assign to first matching picture
            }
        }
    }

    // Decide per-Picture region: suppress or preserve.
    // Empty-validated regions never suppress. Regions with substantive text
    // (>= threshold alphanumeric chars) are preserved as unassigned.
    let mut picture_preserved_segments: std::collections::HashSet<usize> = std::collections::HashSet::new();
    let mut picture_suppressed_segments: std::collections::HashSet<usize> = std::collections::HashSet::new();

    for (pi, seg_indices) in picture_seg_indices.iter().enumerate() {
        let (_, is_empty) = picture_hints[pi];

        if is_empty {
            // Empty-validated: no real image content, don't suppress text
            for &idx in seg_indices {
                picture_preserved_segments.insert(idx);
            }
            continue;
        }

        // Count alphanumeric chars across all segments in this Picture region
        let alphanum_count: usize = seg_indices
            .iter()
            .map(|&idx| segments[idx].text.chars().filter(|c| c.is_alphanumeric()).count())
            .sum();

        if alphanum_count >= PICTURE_SUBSTANTIVE_CHAR_THRESHOLD {
            // Substantive text — likely a screenshot or misclassified region.
            // Preserve segments as unassigned so they go through the standard pipeline.
            tracing::trace!(
                alphanum_count,
                segment_count = seg_indices.len(),
                "picture region contains substantive text — preserving"
            );
            for &idx in seg_indices {
                picture_preserved_segments.insert(idx);
            }
        } else {
            // Short/label text — suppress as before (diagram labels, axis text, etc.)
            for &idx in seg_indices {
                picture_suppressed_segments.insert(idx);
            }
        }
    }

    for (seg_idx, seg) in segments.iter().enumerate() {
        if seg.text.trim().is_empty() {
            continue; // Skip whitespace-only segments
        }

        let seg_rect = Rect::from_lbrt(seg.x, seg.y, seg.x + seg.width, seg.y + seg.height);

        // Suppress segments overlapping successfully extracted tables (>=50% IoS).
        let in_extracted_table = suppress_bboxes.iter().any(|bb| {
            let bb_rect = Rect::from_lbrt(bb.x0 as f32, bb.y0 as f32, bb.x1 as f32, bb.y1 as f32);
            seg_rect.intersection_over_self(&bb_rect) >= 0.5
        });
        if in_extracted_table {
            suppressed_count += 1;
            continue;
        }

        // Handle Picture region segments based on per-region decision
        if picture_suppressed_segments.contains(&seg_idx) {
            suppressed_count += 1;
            continue;
        }
        if picture_preserved_segments.contains(&seg_idx) {
            // Preserved picture text goes to unassigned for standard pipeline processing
            unassigned.push(seg_idx);
            continue;
        }

        let mut best_hint_idx: Option<usize> = None;
        let mut best_ios = 0.0_f32;
        let mut best_area = f32::MAX;

        for (hi, hint) in confident_hints.iter().enumerate() {
            let hint_rect = Rect::from_lbrt(hint.left, hint.bottom, hint.right, hint.top);
            let ios = seg_rect.intersection_over_self(&hint_rect);

            if ios >= MIN_IOS_THRESHOLD {
                // Prefer higher IoS; break ties by smallest area
                if ios > best_ios || (ios == best_ios && hint_areas[hi] < best_area) {
                    best_ios = ios;
                    best_area = hint_areas[hi];
                    best_hint_idx = Some(hi);
                }
            }
        }

        match best_hint_idx {
            Some(hi) => regions[hi].segment_indices.push(seg_idx),
            None => unassigned.push(seg_idx),
        }
    }

    tracing::trace!(
        confident_hints = confident_hints.len(),
        segments = segments.len(),
        suppressed = suppressed_count,
        assigned_to_regions = regions.iter().map(|r| r.segment_indices.len()).sum::<usize>(),
        unassigned = unassigned.len(),
        "segment-to-region assignment complete"
    );

    (regions, unassigned)
}

/// Assign segments to regions with iterative bounding box refinement.
///
/// After the initial assignment, each region's bbox is shrunk to tightly fit its
/// assigned segments (plus padding). Then segments are re-assigned using the refined
/// bboxes. This prevents over-large model bboxes from "stealing" text from adjacent
/// regions. Repeats up to 3 iterations or until assignments stabilize.
pub(in crate::pdf::markdown) fn assign_segments_to_regions_refined<'a>(
    segments: &[SegmentData],
    hints: &'a [LayoutHint],
    min_confidence: f32,
    extracted_table_bboxes: &[crate::types::BoundingBox],
    hint_validations: &[super::layout_validation::RegionValidation],
) -> (Vec<LayoutRegion<'a>>, Vec<usize>) {
    // First pass: assign with original bboxes, using validation for Picture suppression
    let (regions, unassigned) = assign_segments_to_regions(
        segments,
        hints,
        min_confidence,
        extracted_table_bboxes,
        hint_validations,
    );

    if regions.is_empty() {
        return (regions, unassigned);
    }

    // Compute refined bboxes from assigned segments
    let refined_hints: Vec<LayoutHint> = compute_refined_hints(&regions, segments, hints);

    if refined_hints.is_empty() {
        return (regions, unassigned);
    }

    // Collect Picture hints from the original hints to carry through iterations.
    // The refined hints contain text-bearing and Table regions (non-Picture),
    // but Picture bboxes must remain available for segment suppression.
    let picture_hints: Vec<LayoutHint> = hints
        .iter()
        .filter(|h| h.confidence >= min_confidence && h.class == LayoutHintClass::Picture)
        .cloned()
        .collect();

    // Iterate: re-assign with refined bboxes, re-shrink, repeat
    let mut current_hints = refined_hints;
    // Append Picture hints so they're available for suppression in each iteration
    current_hints.extend(picture_hints.iter().cloned());
    let mut prev_assignments: Vec<Vec<usize>> = regions.iter().map(|r| r.segment_indices.clone()).collect();

    for _ in 1..MAX_REFINEMENT_ITERATIONS {
        let (new_regions, _) =
            assign_segments_to_regions(segments, &current_hints, min_confidence, extracted_table_bboxes, &[]);

        // Check if assignments changed
        let new_assignments: Vec<Vec<usize>> = new_regions.iter().map(|r| r.segment_indices.clone()).collect();
        if new_assignments == prev_assignments {
            break; // Stable — stop iterating
        }

        prev_assignments = new_assignments;
        let mut new_refined = compute_refined_hints(&new_regions, segments, &current_hints);
        new_refined.extend(picture_hints.iter().cloned());
        current_hints = new_refined;
    }

    // Final assignment with the last refined hints, but we need to return regions
    // that reference the original hints (for class/confidence), with the refined assignments
    let (final_regions_refined, final_unassigned) =
        assign_segments_to_regions(segments, &current_hints, min_confidence, extracted_table_bboxes, &[]);

    // Map refined regions back to original hints by matching class + position
    let mut result_regions: Vec<LayoutRegion<'a>> = Vec::new();
    let confident_original: Vec<(usize, &'a LayoutHint)> = hints
        .iter()
        .enumerate()
        .filter(|(_, h)| h.confidence >= min_confidence)
        .filter(|(_, h)| !matches!(h.class, LayoutHintClass::Table | LayoutHintClass::Picture))
        .collect();

    for (ri, refined_region) in final_regions_refined.iter().enumerate() {
        if ri < confident_original.len() {
            result_regions.push(LayoutRegion {
                hint: confident_original[ri].1,
                segment_indices: refined_region.segment_indices.clone(),
                merged_bbox: None,
            });
        }
    }

    (result_regions, final_unassigned)
}

/// Compute refined hints by shrinking each region's bbox to tightly fit its segments.
///
/// Each region's `hint` field already points to the correct original hint for that
/// region (set when the region was created in `assign_segments_to_regions`), so we
/// use it directly as the base for clamping. The `source_hints` parameter is unused
/// but retained for API symmetry with the iterative caller.
fn compute_refined_hints(
    regions: &[LayoutRegion],
    segments: &[SegmentData],
    _source_hints: &[LayoutHint],
) -> Vec<LayoutHint> {
    let mut refined = Vec::with_capacity(regions.len());

    for region in regions.iter() {
        let base_hint = region.hint;

        if region.segment_indices.is_empty() {
            // No segments — keep original bbox
            refined.push(base_hint.clone());
            continue;
        }

        // Compute tight bbox from assigned segments
        let mut tight_left = f32::MAX;
        let mut tight_bottom = f32::MAX;
        let mut tight_right = f32::MIN;
        let mut tight_top = f32::MIN;

        for &idx in &region.segment_indices {
            let seg = &segments[idx];
            tight_left = tight_left.min(seg.x);
            tight_bottom = tight_bottom.min(seg.y);
            tight_right = tight_right.max(seg.x + seg.width);
            tight_top = tight_top.max(seg.y + seg.height);
        }

        // Sparsity guard: if segments fill < 15% of the tight bbox area,
        // the layout is sparse (e.g., TOC with dot leaders). Refinement
        // would fragment the line across regions — keep original bbox.
        let tight_area = (tight_right - tight_left) * (tight_top - tight_bottom);
        if tight_area > 0.0 {
            let seg_area_sum: f32 = region
                .segment_indices
                .iter()
                .map(|&idx| segments[idx].width * segments[idx].height)
                .sum();
            let fill_ratio = seg_area_sum / tight_area;
            if fill_ratio < 0.15 {
                tracing::trace!(
                    class = ?base_hint.class,
                    fill_ratio,
                    "refinement skipped: sparse layout"
                );
                refined.push(base_hint.clone());
                continue;
            }
        }

        // Apply padding and clamp to original bbox (never expand beyond model prediction)
        refined.push(LayoutHint {
            class: base_hint.class,
            confidence: base_hint.confidence,
            left: (tight_left - BBOX_REFINEMENT_PADDING).max(base_hint.left),
            bottom: (tight_bottom - BBOX_REFINEMENT_PADDING).max(base_hint.bottom),
            right: (tight_right + BBOX_REFINEMENT_PADDING).min(base_hint.right),
            top: (tight_top + BBOX_REFINEMENT_PADDING).min(base_hint.top),
        });
    }

    refined
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::pdf::hierarchy::SegmentData;
    use crate::pdf::markdown::types::{LayoutHint, LayoutHintClass};

    fn make_segment(text: &str, x: f32, y: f32, w: f32, h: f32) -> SegmentData {
        SegmentData {
            text: text.to_string(),
            x,
            y,
            width: w,
            height: h,
            font_size: 12.0,
            is_bold: false,
            is_italic: false,
            is_monospace: false,
            baseline_y: y,
        }
    }

    fn make_hint(class: LayoutHintClass, left: f32, bottom: f32, right: f32, top: f32) -> LayoutHint {
        LayoutHint {
            class,
            confidence: 0.9,
            left,
            bottom,
            right,
            top,
        }
    }

    #[test]
    fn picture_suppresses_short_label_text() {
        // A Picture region with only short label text (< 50 alphanum chars)
        // should suppress the segments.
        let segments = vec![
            make_segment("Fig 1", 10.0, 10.0, 30.0, 10.0),
            make_segment("x-axis", 10.0, 5.0, 30.0, 10.0),
        ];
        let hints = vec![make_hint(LayoutHintClass::Picture, 0.0, 0.0, 100.0, 100.0)];
        let (regions, unassigned) = assign_segments_to_regions(&segments, &hints, 0.5, &[], &[]);
        // No text regions, no unassigned — both segments suppressed
        assert!(regions.is_empty());
        assert!(
            unassigned.is_empty(),
            "short label text should be suppressed, got {:?}",
            unassigned
        );
    }

    #[test]
    fn picture_preserves_substantive_text() {
        // A Picture region containing substantial readable text (e.g., a screenshot
        // of a paper) should preserve segments as unassigned.
        let long_text =
            "This is a substantial amount of readable text that should not be suppressed by the layout model";
        let segments = vec![
            make_segment(long_text, 10.0, 50.0, 200.0, 12.0),
            make_segment(
                "Additional paragraph of text in the screenshot region",
                10.0,
                30.0,
                200.0,
                12.0,
            ),
        ];
        let hints = vec![make_hint(LayoutHintClass::Picture, 0.0, 0.0, 300.0, 100.0)];
        let (regions, unassigned) = assign_segments_to_regions(&segments, &hints, 0.5, &[], &[]);
        // Both segments should be preserved as unassigned
        assert!(regions.is_empty());
        assert_eq!(
            unassigned.len(),
            2,
            "substantive text should be preserved as unassigned"
        );
    }

    #[test]
    fn picture_empty_validated_never_suppresses() {
        // A Picture region validated as Empty should never suppress text.
        let segments = vec![make_segment("Fig 1", 10.0, 10.0, 30.0, 10.0)];
        let hints = vec![make_hint(LayoutHintClass::Picture, 0.0, 0.0, 100.0, 100.0)];
        let validations = vec![super::super::layout_validation::RegionValidation::Empty];
        let (regions, unassigned) = assign_segments_to_regions(&segments, &hints, 0.5, &[], &validations);
        assert!(regions.is_empty());
        // Even short text should be preserved when region is empty-validated
        assert_eq!(unassigned.len(), 1, "empty-validated picture should not suppress text");
    }

    #[test]
    fn picture_does_not_affect_non_overlapping_segments() {
        // Segments outside the Picture region should be unaffected.
        let segments = vec![
            make_segment("Outside text", 200.0, 200.0, 100.0, 12.0),
            make_segment("Label inside", 10.0, 10.0, 30.0, 10.0),
        ];
        let hints = vec![make_hint(LayoutHintClass::Picture, 0.0, 0.0, 100.0, 100.0)];
        let (regions, unassigned) = assign_segments_to_regions(&segments, &hints, 0.5, &[], &[]);
        // "Outside text" should be unassigned (no non-Picture regions to match)
        // "Label inside" should be suppressed (short label text)
        assert!(regions.is_empty());
        assert_eq!(unassigned.len(), 1);
        assert_eq!(unassigned[0], 0, "only the outside segment should be unassigned");
    }

    #[test]
    fn picture_threshold_boundary() {
        // Test right at the threshold boundary (50 alphanumeric chars).
        // Generate exactly 50 alphanum chars.
        let text_50 = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWX"; // 50 chars
        assert_eq!(text_50.chars().filter(|c| c.is_alphanumeric()).count(), 50);

        let segments = vec![make_segment(text_50, 10.0, 10.0, 200.0, 12.0)];
        let hints = vec![make_hint(LayoutHintClass::Picture, 0.0, 0.0, 300.0, 100.0)];
        let (_, unassigned) = assign_segments_to_regions(&segments, &hints, 0.5, &[], &[]);
        // Exactly at threshold — should be preserved
        assert_eq!(unassigned.len(), 1, "text at threshold should be preserved");

        // Now test with 49 chars (below threshold)
        let text_49 = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVW"; // 49 chars
        assert_eq!(text_49.chars().filter(|c| c.is_alphanumeric()).count(), 49);

        let segments_below = vec![make_segment(text_49, 10.0, 10.0, 200.0, 12.0)];
        let (_, unassigned_below) = assign_segments_to_regions(&segments_below, &hints, 0.5, &[], &[]);
        // Below threshold — should be suppressed
        assert!(unassigned_below.is_empty(), "text below threshold should be suppressed");
    }

    #[test]
    fn mixed_picture_and_text_regions() {
        // Picture region + Text region on the same page.
        // Segments in the Text region should be assigned normally.
        // Segments in the Picture region with short text should be suppressed.
        let segments = vec![
            make_segment("Body paragraph text", 10.0, 200.0, 150.0, 12.0),
            make_segment("Fig 1", 10.0, 50.0, 30.0, 10.0),
        ];
        let hints = vec![
            make_hint(LayoutHintClass::Text, 0.0, 180.0, 200.0, 230.0),
            make_hint(LayoutHintClass::Picture, 0.0, 0.0, 100.0, 100.0),
        ];
        let (regions, unassigned) = assign_segments_to_regions(&segments, &hints, 0.5, &[], &[]);
        // "Body paragraph text" should be assigned to the Text region
        assert_eq!(regions.len(), 1);
        assert_eq!(regions[0].hint.class, LayoutHintClass::Text);
        assert_eq!(regions[0].segment_indices, vec![0]);
        // "Fig 1" should be suppressed (short label)
        assert!(unassigned.is_empty());
    }
}