layout_table_detection.py

# src/layout_table_detection.py

"""
Phase 4.x: Table Detection via Grid Pattern Analysis

This module detects tabular structures in OCR blocks using:
- Simple mode: Alignment-based grid detection (fast)
- Advanced mode: Density + whitespace analysis (more accurate)
"""

from typing import List, Dict, Any
from layout_blocks import OCRBlock


def detect_tables_simple(
    blocks: List[OCRBlock],
    min_grid_score: float = 0.60,
    min_rows: int = 3,
    min_cols: int = 2,
) -> List[OCRBlock]:
    """
    Detect tables using alignment heuristics (simple mode).

    Algorithm:
    1. Find dense rectangular regions (min 3 rows × 2 cols)
    2. Cluster blocks into rows (y-coordinate proximity)
    3. Cluster blocks into columns (x-coordinate proximity)
    4. Calculate grid_score = (row_alignment + col_alignment) / 2
    5. Mark blocks as tables if grid_score ≥ threshold

    Args:
        blocks: OCR blocks for a page
        min_grid_score: Minimum alignment score (0.0-1.0, default: 0.60)
        min_rows: Minimum rows for table detection (default: 3)
        min_cols: Minimum columns for table detection (default: 2)

    Returns:
        Blocks with table_metadata populated for detected tables

    Performance: O(n²), target <50ms per page
    """
    if not blocks or len(blocks) < min_rows * min_cols:
        return blocks

    # Step 1: Find dense regions (candidates for tables)
    candidates = _find_dense_regions(blocks, min_blocks=min_rows * min_cols)

    table_id = 0

    for region in candidates:
        # Step 2: Cluster into rows (y-coordinate)
        rows = _cluster_by_y_coordinate(region, tolerance=0.02)

        if len(rows) < min_rows:
            continue

        # Step 3: Cluster into columns (x-coordinate)
        cols = _cluster_by_x_coordinate(region, tolerance=0.05)

        if len(cols) < min_cols:
            continue

        # Step 4: Calculate grid alignment scores
        row_score = _calculate_row_alignment_score(rows)
        col_score = _calculate_col_alignment_score(cols)
        grid_score = (row_score + col_score) / 2.0

        if grid_score >= min_grid_score:
            # Step 5: Mark blocks as table
            for block in region:
                block.table_metadata = {
                    "is_table": True,
                    "table_id": table_id,
                    "detection_mode": "simple",
                    "grid_score": round(grid_score, 3),
                    "rows": len(rows),
                    "cols": len(cols),
                    "confidence": round(grid_score, 3),
                }
            table_id += 1

    return blocks


def _find_dense_regions(
    blocks: List[OCRBlock],
    min_blocks: int,
    density_threshold: float = 0.25,
) -> List[List[OCRBlock]]:
    """
    Find rectangular regions with high block density.

    Uses spatial clustering to group nearby blocks.

    Args:
        blocks: OCR blocks to analyze
        min_blocks: Minimum blocks required in a region
        density_threshold: Maximum distance for clustering (normalized)

    Returns:
        List of block regions (each region is a list of blocks)
    """
    if not blocks:
        return []

    # DBSCAN-like clustering with proper expansion
    visited = set()
    regions = []

    for i, block in enumerate(blocks):
        if i in visited:
            continue

        # Start new region with BFS/DFS expansion
        region = []
        queue = [i]

        while queue:
            current_idx = queue.pop(0)

            if current_idx in visited:
                continue

            visited.add(current_idx)
            region.append(blocks[current_idx])

            # Find all unvisited neighbors of current block
            for j, other in enumerate(blocks):
                if j not in visited:
                    if _are_blocks_close(blocks[current_idx], other, density_threshold):
                        queue.append(j)

        # Keep regions with enough blocks
        if len(region) >= min_blocks:
            regions.append(region)

    return regions


def _are_blocks_close(block1: OCRBlock, block2: OCRBlock, threshold: float) -> bool:
    """
    Check if two blocks are spatially close.

    Uses center-to-center distance with normalized coordinates.
    """
    dx = block1.x_center - block2.x_center
    dy = block1.y_center - block2.y_center
    distance = (dx ** 2 + dy ** 2) ** 0.5
    return distance <= threshold  # Use <= to include boundary cases


def _cluster_by_y_coordinate(
    blocks: List[OCRBlock],
    tolerance: float,
) -> List[List[OCRBlock]]:
    """
    Group blocks into rows based on y-coordinate proximity.

    Args:
        blocks: Blocks to cluster
        tolerance: Maximum y-distance for same row (normalized)

    Returns:
        List of rows (each row is a list of blocks)
    """
    if not blocks:
        return []

    # Sort by y0 (top coordinate)
    sorted_blocks = sorted(blocks, key=lambda b: b.bbox[1])

    rows = []
    current_row = [sorted_blocks[0]]

    for block in sorted_blocks[1:]:
        prev_y = current_row[-1].bbox[1]  # y0
        curr_y = block.bbox[1]

        if abs(curr_y - prev_y) <= tolerance:
            current_row.append(block)
        else:
            rows.append(current_row)
            current_row = [block]

    rows.append(current_row)
    return rows


def _cluster_by_x_coordinate(
    blocks: List[OCRBlock],
    tolerance: float,
) -> List[List[OCRBlock]]:
    """
    Group blocks into columns based on x-coordinate proximity.

    Args:
        blocks: Blocks to cluster
        tolerance: Maximum x-distance for same column (normalized)

    Returns:
        List of columns (each column is a list of blocks)
    """
    if not blocks:
        return []

    # Sort by x0 (left coordinate)
    sorted_blocks = sorted(blocks, key=lambda b: b.bbox[0])

    cols = []
    current_col = [sorted_blocks[0]]

    for block in sorted_blocks[1:]:
        prev_x = current_col[-1].bbox[0]  # x0
        curr_x = block.bbox[0]

        if abs(curr_x - prev_x) <= tolerance:
            current_col.append(block)
        else:
            cols.append(current_col)
            current_col = [block]

    cols.append(current_col)
    return cols


def _calculate_row_alignment_score(rows: List[List[OCRBlock]]) -> float:
    """
    Calculate how well rows are aligned (uniform block count per row).

    Perfect alignment (all rows have same number of blocks) = 1.0
    High variance in row lengths = lower score

    Args:
        rows: List of rows (each row is list of blocks)

    Returns:
        Alignment score (0.0-1.0)
    """
    if not rows:
        return 0.0

    row_lengths = [len(row) for row in rows]
    avg = sum(row_lengths) / len(row_lengths)

    if avg == 0:
        return 0.0

    # Calculate coefficient of variation
    variance = sum((x - avg) ** 2 for x in row_lengths) / len(row_lengths)
    std_dev = variance ** 0.5
    coefficient_of_variation = std_dev / avg

    # Convert to score: low variance = high score
    score = max(0.0, 1.0 - coefficient_of_variation)

    return score


def _calculate_col_alignment_score(cols: List[List[OCRBlock]]) -> float:
    """
    Calculate how well columns are aligned (uniform block count per column).

    Args:
        cols: List of columns (each column is list of blocks)

    Returns:
        Alignment score (0.0-1.0)
    """
    if not cols:
        return 0.0

    col_lengths = [len(col) for col in cols]
    avg = sum(col_lengths) / len(col_lengths)

    if avg == 0:
        return 0.0

    # Calculate coefficient of variation
    variance = sum((x - avg) ** 2 for x in col_lengths) / len(col_lengths)
    std_dev = variance ** 0.5
    coefficient_of_variation = std_dev / avg

    # Convert to score
    score = max(0.0, 1.0 - coefficient_of_variation)

    return score


# Advanced mode functions

def detect_tables_advanced(
    blocks: List[OCRBlock],
    min_confidence: float = 0.65,
) -> List[OCRBlock]:
    """
    Advanced table detection with density and whitespace analysis.

    Runs simple mode first (lower threshold), then refines with additional metrics.

    Args:
        blocks: OCR blocks for a page
        min_confidence: Minimum combined confidence (default: 0.65)

    Returns:
        Blocks with refined table_metadata

    Performance: O(n² log n), target <200ms per page
    """
    # Run simple mode with lower threshold
    blocks = detect_tables_simple(blocks, min_grid_score=0.50)

    # Get unique table IDs
    table_ids = set()
    for block in blocks:
        if block.table_metadata and block.table_metadata["is_table"]:
            table_ids.add(block.table_metadata["table_id"])

    # Refine each detected table
    for table_id in table_ids:
        table_blocks = [
            b for b in blocks
            if b.table_metadata and b.table_metadata.get("table_id") == table_id
        ]

        if not table_blocks:
            continue

        # Calculate additional metrics
        density_score = _calculate_density_score(table_blocks)
        whitespace_score = _calculate_whitespace_uniformity(table_blocks)
        alignment_score = table_blocks[0].table_metadata["grid_score"]

        # Combined confidence: 40% alignment + 30% density + 30% whitespace
        confidence = (
            0.4 * alignment_score +
            0.3 * density_score +
            0.3 * whitespace_score
        )

        if confidence >= min_confidence:
            # Update metadata with advanced mode results
            for block in table_blocks:
                block.table_metadata.update({
                    "detection_mode": "advanced",
                    "confidence": round(confidence, 3),
                    "density_score": round(density_score, 3),
                    "whitespace_score": round(whitespace_score, 3),
                })
        else:
            # Remove table metadata (false positive)
            for block in table_blocks:
                block.table_metadata = None

    return blocks


def _calculate_density_score(table_blocks: List[OCRBlock]) -> float:
    """
    Calculate block-area density for a table region.

    Higher density indicates tighter packing typical of tables.

    Args:
        table_blocks: Blocks in the table region

    Returns:
        Density score (0.0-1.0)
    """
    if not table_blocks:
        return 0.0

    # Calculate bounding box of entire table
    x0 = min(b.bbox[0] for b in table_blocks)
    y0 = min(b.bbox[1] for b in table_blocks)
    x1 = max(b.bbox[2] for b in table_blocks)
    y1 = max(b.bbox[3] for b in table_blocks)

    table_area = (x1 - x0) * (y1 - y0)

    if table_area == 0:
        return 0.0

    # Calculate total area of blocks
    blocks_area = sum(b.area for b in table_blocks)

    # Density = blocks_area / table_area
    density = blocks_area / table_area

    # Normalize to 0-1 (tables typically have 0.2-0.6 density)
    # Map 0.2-0.6 to 0.5-1.0
    if density < 0.2:
        return density * 2.5  # 0.2 → 0.5
    elif density > 0.6:
        return 1.0
    else:
        return 0.5 + (density - 0.2) * 1.25  # 0.2-0.6 → 0.5-1.0

    return min(1.0, density)


def _calculate_whitespace_uniformity(table_blocks: List[OCRBlock]) -> float:
    """
    Calculate uniformity of inter-cell gaps (whitespace).

    Tables have consistent spacing; narrative text has variable spacing.

    Args:
        table_blocks: Blocks in the table region

    Returns:
        Uniformity score (0.0-1.0)
    """
    if len(table_blocks) < 2:
        return 0.5  # Neutral score for single-block "tables"

    # Calculate horizontal gaps between adjacent blocks (same row)
    h_gaps = []
    # Calculate vertical gaps between adjacent blocks (same column)
    v_gaps = []

    for i, block in enumerate(table_blocks):
        for other in table_blocks[i + 1:]:
            # Check if roughly same y-position (horizontal neighbors)
            if abs(block.y_center - other.y_center) < 0.05:
                gap = abs(block.bbox[2] - other.bbox[0])  # x1 of one, x0 of other
                if gap < 0.3:  # Reasonable max gap
                    h_gaps.append(gap)

            # Check if roughly same x-position (vertical neighbors)
            if abs(block.x_center - other.x_center) < 0.05:
                gap = abs(block.bbox[3] - other.bbox[1])  # y1 of one, y0 of other
                if gap < 0.3:
                    v_gaps.append(gap)

    # Calculate uniformity (low std dev = high uniformity)
    all_gaps = h_gaps + v_gaps

    if not all_gaps:
        return 0.5

    avg_gap = sum(all_gaps) / len(all_gaps)

    if avg_gap == 0:
        return 1.0  # Perfect uniformity (no gaps)

    variance = sum((g - avg_gap) ** 2 for g in all_gaps) / len(all_gaps)
    std_dev = variance ** 0.5
    coefficient_of_variation = std_dev / avg_gap

    # Low CV = high uniformity
    score = max(0.0, 1.0 - coefficient_of_variation)

    return score


# Helper functions

def is_table_block(block: OCRBlock) -> bool:
    """Check if a block is marked as part of a table."""
    return (
        block.table_metadata is not None and
        block.table_metadata.get("is_table", False)
    )


def get_table_blocks(blocks: List[OCRBlock], table_id: int) -> List[OCRBlock]:
    """Get all blocks belonging to a specific table."""
    return [
        b for b in blocks
        if b.table_metadata and b.table_metadata.get("table_id") == table_id
    ]


def count_tables(blocks: List[OCRBlock]) -> int:
    """Count the number of detected tables on a page."""
    table_ids = set()
    for block in blocks:
        if is_table_block(block):
            table_ids.add(block.table_metadata["table_id"])
    return len(table_ids)