layout_grouping.py

# src/layout_grouping.py

"""
Phase 4.x: Paragraph Grouping via Spatial Cues

This module groups OCR blocks into paragraphs using vertical gap analysis,
replacing text-based paragraph detection with spatial structure preservation.
"""

from typing import List
from layout_blocks import OCRBlock


def group_paragraphs(
    blocks: List[OCRBlock],
    gap_multiplier: float = 1.5,
    min_gap_threshold: float = 0.01,
) -> List[OCRBlock]:
    """
    Assign paragraph IDs to blocks based on vertical gaps between consecutive blocks.

    Algorithm:
    1. Calculate vertical gaps between consecutive blocks (y1 of block N → y0 of block N+1)
    2. Compute dynamic threshold: median_gap × gap_multiplier
    3. Assign sequential paragraph_id to blocks, starting new paragraph when gap > threshold

    Args:
        blocks: OCR blocks in reading order (must be pre-sorted)
        gap_multiplier: Multiplier for median gap to determine paragraph breaks (default: 1.5)
        min_gap_threshold: Minimum gap to consider as paragraph break (normalized, default: 0.01)

    Returns:
        List of OCRBlock with paragraph_id assigned (modifies blocks in-place)

    Performance:
        O(n) for gap calculation + O(n log n) for median = O(n log n) overall
        Target: <5ms per page with 100 blocks
    """
    if not blocks:
        return []

    # Edge case: single block
    if len(blocks) == 1:
        blocks[0].paragraph_id = 0
        return blocks

    # Step 1: Calculate vertical gaps between consecutive blocks
    gaps = []
    for i in range(len(blocks) - 1):
        current_bottom = blocks[i].bbox[3]  # y1
        next_top = blocks[i + 1].bbox[1]    # y0
        gap = next_top - current_bottom
        gaps.append(gap)

    # Step 2: Compute dynamic threshold based on minimum gap
    # Strategy: Use the smallest gap as baseline for "typical within-paragraph spacing"
    # Any gap significantly larger (gap_multiplier × min_gap) is a paragraph break
    # This is robust because within-paragraph gaps are typically consistent and small

    if len(gaps) == 1:
        # Single gap case: compare against min_gap_threshold
        threshold = min_gap_threshold
    else:
        sorted_gaps = sorted(gaps)
        # Use minimum gap as baseline for typical within-paragraph spacing
        min_actual_gap = sorted_gaps[0]
        # Threshold is minimum gap × multiplier, capped by min_gap_threshold
        threshold = max(min_actual_gap * gap_multiplier, min_gap_threshold)

    # Step 3: Assign paragraph IDs based on threshold
    paragraph_id = 0
    blocks[0].paragraph_id = paragraph_id

    for i in range(1, len(blocks)):
        gap = gaps[i - 1]
        if gap > threshold:
            paragraph_id += 1
        blocks[i].paragraph_id = paragraph_id

    return blocks


def get_paragraph_blocks(
    blocks: List[OCRBlock],
    paragraph_id: int,
) -> List[OCRBlock]:
    """
    Retrieve all blocks belonging to a specific paragraph.

    Args:
        blocks: List of OCRBlock instances with paragraph_id assigned
        paragraph_id: Target paragraph ID to retrieve

    Returns:
        List of blocks belonging to the specified paragraph
    """
    return [b for b in blocks if b.paragraph_id == paragraph_id]


def count_paragraphs(blocks: List[OCRBlock]) -> int:
    """
    Count the total number of paragraphs in a block list.

    Args:
        blocks: List of OCRBlock instances with paragraph_id assigned

    Returns:
        Number of unique paragraphs (0 if no blocks or no IDs assigned)
    """
    if not blocks:
        return 0

    # Get maximum paragraph_id (paragraphs are 0-indexed)
    max_id = max(
        (b.paragraph_id for b in blocks if b.paragraph_id is not None),
        default=None
    )

    return 0 if max_id is None else max_id + 1


def get_paragraph_text(
    blocks: List[OCRBlock],
    paragraph_id: int,
    separator: str = " ",
) -> str:
    """
    Extract concatenated text from a specific paragraph.

    Args:
        blocks: List of OCRBlock instances
        paragraph_id: Target paragraph ID
        separator: String to join block texts (default: space)

    Returns:
        Combined text from all blocks in the paragraph
    """
    paragraph_blocks = get_paragraph_blocks(blocks, paragraph_id)
    return separator.join(b.text for b in paragraph_blocks)


def get_paragraph_stats(blocks: List[OCRBlock]) -> dict:
    """
    Calculate statistics about paragraph distribution.

    Args:
        blocks: List of OCRBlock instances with paragraph_id assigned

    Returns:
        Dictionary with paragraph statistics:
        - total_paragraphs: Number of paragraphs
        - avg_blocks_per_paragraph: Average blocks per paragraph
        - min_blocks: Minimum blocks in any paragraph
        - max_blocks: Maximum blocks in any paragraph
    """
    if not blocks:
        return {
            "total_paragraphs": 0,
            "avg_blocks_per_paragraph": 0.0,
            "min_blocks": 0,
            "max_blocks": 0,
        }

    total_paragraphs = count_paragraphs(blocks)

    # Count blocks per paragraph
    paragraph_counts = {}
    for block in blocks:
        if block.paragraph_id is not None:
            paragraph_counts[block.paragraph_id] = (
                paragraph_counts.get(block.paragraph_id, 0) + 1
            )

    counts = list(paragraph_counts.values())

    return {
        "total_paragraphs": total_paragraphs,
        "avg_blocks_per_paragraph": sum(counts) / len(counts) if counts else 0.0,
        "min_blocks": min(counts) if counts else 0,
        "max_blocks": max(counts) if counts else 0,
    }