layout_order.py

"""
Reading order resolution using X-histogram column detection.

This module provides robust multi-column detection and reading order
resolution for OCR text blocks using spatial analysis.
"""

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


# ---------------------------------------
# X-Histogram Helper Functions
# ---------------------------------------

def _moving_average(data: List[float], window: int) -> List[float]:
    """
    Smooth histogram with moving average filter.

    Args:
        data: Input histogram values
        window: Window size for averaging (must be odd for symmetry)

    Returns:
        Smoothed histogram values
    """
    if window < 1:
        return data

    if len(data) < window:
        return data

    smoothed = []
    half_window = window // 2

    for i in range(len(data)):
        # Compute window bounds
        start = max(0, i - half_window)
        end = min(len(data), i + half_window + 1)

        # Average values in window
        avg = sum(data[start:end]) / (end - start)
        smoothed.append(avg)

    return smoothed


def _find_valleys(histogram: List[float], min_depth: float = 0.2) -> List[int]:
    """
    Find local minima in histogram as column boundaries.

    A valley is a local minimum where the value is significantly
    lower than surrounding peaks.

    Args:
        histogram: Histogram values
        min_depth: Minimum depth ratio (valley must be at least
                   this fraction below neighboring peaks)

    Returns:
        List of bin indices representing valley positions
    """
    if len(histogram) < 3:
        return []

    # Find global max for depth calculation
    max_val = max(histogram) if histogram else 1.0
    if max_val == 0:
        return []

    valleys = []

    for i in range(1, len(histogram) - 1):
        left = histogram[i - 1]
        center = histogram[i]
        right = histogram[i + 1]

        # Check if local minimum
        if center < left and center < right:
            # Check depth requirement
            peak = max(left, right)
            depth = (peak - center) / max_val if max_val > 0 else 0

            if depth >= min_depth:
                valleys.append(i)

    return valleys


# ---------------------------------------
# Column Detection
# ---------------------------------------

def detect_columns(
    blocks: List[OCRBlock],
    resolution: int = 100,
    smooth_window: int = 5,
    min_depth: float = 0.2,
    max_columns: int = 4,
) -> List[List[OCRBlock]]:
    """
    Detect columns using X-coordinate histogram clustering.

    Algorithm:
    1. Build X-histogram: count blocks per bin across page width
    2. Smooth histogram with moving average
    3. Find valleys (local minima) as column boundaries
    4. Assign blocks to columns based on x_center
    5. Sort columns left-to-right

    Args:
        blocks: OCR blocks to cluster
        resolution: Number of histogram bins (default: 100)
        smooth_window: Moving average window size (default: 5)
        min_depth: Minimum valley depth ratio (default: 0.2)
        max_columns: Maximum columns to detect (default: 4)

    Returns:
        List of columns, each containing blocks
    """
    if not blocks:
        return []

    # Build X-histogram
    histogram = [0.0] * resolution

    for block in blocks:
        # x_center is normalized [0, 1]
        bin_idx = int(block.x_center * (resolution - 1))
        bin_idx = max(0, min(resolution - 1, bin_idx))
        histogram[bin_idx] += 1

    # Smooth histogram
    smoothed = _moving_average(histogram, smooth_window)

    # Find valley positions
    valley_bins = _find_valleys(smoothed, min_depth)

    # Convert bin indices to normalized x-coordinates (boundaries)
    boundaries = [v / (resolution - 1) for v in valley_bins]

    # Limit number of columns
    if len(boundaries) >= max_columns:
        # Keep only deepest valleys
        valley_depths = []
        for v in valley_bins:
            left = smoothed[v - 1] if v > 0 else 0
            right = smoothed[v + 1] if v < len(smoothed) - 1 else 0
            peak = max(left, right)
            depth = peak - smoothed[v]
            valley_depths.append((depth, v))

        # Sort by depth and keep top (max_columns - 1)
        valley_depths.sort(reverse=True)
        top_valleys = sorted([v for _, v in valley_depths[:max_columns - 1]])
        boundaries = [v / (resolution - 1) for v in top_valleys]

    # Assign blocks to columns using boundaries
    # boundaries = [0.3, 0.6] means:
    #   column 0: x_center < 0.3
    #   column 1: 0.3 <= x_center < 0.6
    #   column 2: x_center >= 0.6

    num_cols = len(boundaries) + 1
    columns = [[] for _ in range(num_cols)]

    for block in blocks:
        x = block.x_center

        # Find column index
        col_idx = 0
        for boundary in boundaries:
            if x >= boundary:
                col_idx += 1
            else:
                break

        columns[col_idx].append(block)

    # Remove empty columns
    columns = [col for col in columns if col]

    return columns


def detect_columns_with_metadata(
    blocks: List[OCRBlock],
    resolution: int = 100,
) -> Dict:
    """
    Detect columns and return detailed metadata for debugging.

    Args:
        blocks: OCR blocks to analyze
        resolution: Histogram resolution

    Returns:
        Dictionary with:
        - columns: List of column blocks
        - boundaries: List of boundary x-coordinates
        - histogram: Raw histogram values
        - smoothed: Smoothed histogram values
    """
    if not blocks:
        return {
            "columns": [],
            "boundaries": [],
            "histogram": [],
            "smoothed": [],
        }

    # Build histogram
    histogram = [0.0] * resolution
    for block in blocks:
        bin_idx = int(block.x_center * (resolution - 1))
        bin_idx = max(0, min(resolution - 1, bin_idx))
        histogram[bin_idx] += 1

    # Smooth
    smoothed = _moving_average(histogram, window=5)

    # Valleys
    valley_bins = _find_valleys(smoothed, min_depth=0.2)
    boundaries = [v / (resolution - 1) for v in valley_bins]

    # Detect columns
    columns = detect_columns(blocks, resolution=resolution)

    return {
        "columns": columns,
        "boundaries": boundaries,
        "histogram": histogram,
        "smoothed": smoothed,
    }


# ---------------------------------------
# Reading Order Resolution
# ---------------------------------------

def resolve_reading_order(blocks: List[OCRBlock]) -> List[OCRBlock]:
    """
    Resolve reading order for a single page using column detection.

    Algorithm:
    1. Detect columns via X-histogram
    2. Sort columns left-to-right (by median x_center)
    3. Within each column, sort top-to-bottom (y0), then left-to-right (x0)
    4. Concatenate columns in reading order

    Args:
        blocks: OCR blocks for a single page

    Returns:
        Blocks in correct reading order
    """
    if not blocks:
        return []

    # Detect columns
    columns = detect_columns(blocks)

    if not columns:
        return []

    # Sort columns left-to-right by median x_center
    def column_x_position(col: List[OCRBlock]) -> float:
        if not col:
            return 0.0
        x_centers = [b.x_center for b in col]
        return sorted(x_centers)[len(x_centers) // 2]

    columns.sort(key=column_x_position)

    # Sort blocks within each column top-to-bottom, left-to-right
    ordered = []
    for column in columns:
        # Sort by y0 (top), then x0 (left)
        column_sorted = sorted(column, key=lambda b: (b.bbox[1], b.bbox[0]))
        ordered.extend(column_sorted)

    return ordered

# ---------------------------------------
# Phase 4.x: Sidebar Detection
# ---------------------------------------

def detect_sidebars(
    blocks: List[OCRBlock],
    narrow_threshold: float = 0.35,
    margin_threshold: float = 0.15,
) -> List[OCRBlock]:
    """
    Detect sidebar and callout content based on column width and position.

    Args:
        blocks: OCR blocks for a page
        narrow_threshold: Maximum width ratio for sidebars (default: 0.35 = 35% of page)
        margin_threshold: Distance from edge for margin sidebars (default: 0.15)

    Returns:
        Blocks with sidebar_metadata populated

    Performance: O(n) leveraging existing column detection
    """
    if not blocks:
        return blocks

    # Detect columns
    columns = detect_columns(blocks)

    for column in columns:
        if not column:
            continue

        # Calculate column bounding box
        col_x_min = min(b.bbox[0] for b in column)
        col_x_max = max(b.bbox[2] for b in column)
        col_width = col_x_max - col_x_min

        # Check if narrow
        is_narrow = col_width < narrow_threshold

        # Check position
        is_left_margin = col_x_min < margin_threshold
        is_right_margin = col_x_max > (1.0 - margin_threshold)

        # Determine sidebar type
        sidebar_type = None
        position = None

        if is_narrow and is_left_margin:
            sidebar_type = "left"
            position = "left_margin"
        elif is_narrow and is_right_margin:
            sidebar_type = "right"
            position = "right_margin"
        elif is_narrow and not is_left_margin and not is_right_margin:
            sidebar_type = "inset"
            position = "embedded"

        # Mark blocks
        if sidebar_type:
            for block in column:
                block.sidebar_metadata = {
                    "is_sidebar": True,
                    "sidebar_type": sidebar_type,
                    "width_ratio": round(col_width, 3),
                    "position": position,
                }

    return blocks


def is_sidebar_block(block: OCRBlock) -> bool:
    """Check if a block is marked as sidebar content."""
    return (
        block.sidebar_metadata is not None and
        block.sidebar_metadata.get("is_sidebar", False)
    )


def filter_sidebars(blocks: List[OCRBlock]) -> List[OCRBlock]:
    """Remove sidebar blocks from main content flow."""
    return [b for b in blocks if not is_sidebar_block(b)]


def get_sidebar_blocks(blocks: List[OCRBlock]) -> List[OCRBlock]:
    """Get only sidebar blocks."""
    return [b for b in blocks if is_sidebar_block(b)]