visual_chunk_linker.py

"""
Visual chunk linker for Phase 5.v Vision-RAG.

This module links visual ROIs (extracted figures/tables) to semantic chunks
based on page overlap and spatial proximity.
"""

from pathlib import Path
from typing import List, Dict, Any
import json

from logger import get_logger

logger = get_logger()


def link_rois_to_chunks(
    chunks: List[Dict[str, Any]],
    rois: List[Dict[str, Any]],
    proximity_mode: str = "page_overlap"
) -> List[Dict[str, Any]]:
    """
    Link visual ROIs to semantic chunks based on spatial proximity.

    Args:
        chunks: List of semantic chunk dictionaries
        rois: List of ROI metadata from VisualROIExtractor
        proximity_mode: Linking strategy:
            - "page_overlap": Link ROIs to chunks on same page
            - "strict_page": Only link if ROI fully within chunk's page range

    Returns:
        Chunks with updated metadata including visual content references

    Algorithm:
        1. For each ROI:
           - Find chunks that overlap with ROI's page
           - Prioritize chunks on exact page match
           - Add visual metadata to matching chunks
        2. For chunks with multiple ROIs:
           - Keep all references (list of ROI paths)
        3. Set has_visual_content flag for chunks with ROIs
    """
    logger.info(f"Linking {len(rois)} ROIs to {len(chunks)} chunks...")

    # Build page → ROI mapping for fast lookup
    page_to_rois = {}
    for roi in rois:
        page = roi["page"]
        if page not in page_to_rois:
            page_to_rois[page] = []
        page_to_rois[page].append(roi)

    # Link ROIs to chunks
    linked_count = 0
    chunks_with_visuals = 0

    for chunk in chunks:
        metadata = chunk.get("metadata", {})
        page_start = metadata.get("page_start", 0)
        page_end = metadata.get("page_end", 0)

        # Find ROIs that overlap with this chunk's page range
        overlapping_rois = []

        for page in range(page_start, page_end + 1):
            if page in page_to_rois:
                overlapping_rois.extend(page_to_rois[page])

        # Add visual metadata if ROIs found
        if overlapping_rois:
            # Take the first ROI as primary (could be enhanced with better selection)
            primary_roi = overlapping_rois[0]

            metadata["has_visual_content"] = True
            metadata["image_path"] = primary_roi["image_path"]
            metadata["visual_type"] = primary_roi["visual_type"]
            metadata["bbox"] = primary_roi["bbox"]
            metadata["roi_id"] = primary_roi["roi_id"]

            # Store all ROI paths if multiple
            if len(overlapping_rois) > 1:
                metadata["all_image_paths"] = [roi["image_path"] for roi in overlapping_rois]
                metadata["all_visual_types"] = [roi["visual_type"] for roi in overlapping_rois]

            chunk["metadata"] = metadata
            chunks_with_visuals += 1
            linked_count += len(overlapping_rois)

    logger.info(
        f"Linked {linked_count} ROIs to {chunks_with_visuals}/{len(chunks)} chunks "
        f"({chunks_with_visuals/len(chunks)*100:.1f}%)"
    )

    return chunks


def add_visual_metadata_to_chunks_file(
    chunks_path: Path,
    rois: List[Dict[str, Any]],
    output_path: Path = None
) -> Path:
    """
    Add visual metadata to existing chunks JSON file.

    Args:
        chunks_path: Path to chunks_semantic.json
        rois: List of ROI metadata
        output_path: Output path (default: overwrite input)

    Returns:
        Path to updated chunks file

    Example:
        >>> rois = extract_rois_from_pipeline(...)
        >>> add_visual_metadata_to_chunks_file(
        ...     chunks_path=Path("chunks_semantic.json"),
        ...     rois=rois
        ... )
    """
    # Load chunks
    with open(chunks_path, "r", encoding="utf-8") as f:
        chunks = json.load(f)

    # Link ROIs to chunks
    chunks = link_rois_to_chunks(chunks, rois)

    # Save updated chunks
    if output_path is None:
        output_path = chunks_path

    with open(output_path, "w", encoding="utf-8") as f:
        json.dump(chunks, f, indent=2, ensure_ascii=False)

    logger.info(f"Saved chunks with visual metadata to {output_path}")

    return output_path


def get_visual_chunks(chunks: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
    """
    Filter chunks that have visual content.

    Args:
        chunks: List of chunks

    Returns:
        List of chunks with has_visual_content=True
    """
    visual_chunks = [
        chunk for chunk in chunks
        if chunk.get("metadata", {}).get("has_visual_content", False)
    ]

    return visual_chunks


def get_visual_stats(chunks: List[Dict[str, Any]]) -> Dict[str, Any]:
    """
    Get statistics about visual content in chunks.

    Args:
        chunks: List of chunks

    Returns:
        Dictionary with visual content statistics
    """
    visual_chunks = get_visual_chunks(chunks)

    stats = {
        "total_chunks": len(chunks),
        "chunks_with_visuals": len(visual_chunks),
        "visual_percentage": len(visual_chunks) / len(chunks) * 100 if chunks else 0,
        "visual_types": {},
    }

    # Count by visual type
    for chunk in visual_chunks:
        vtype = chunk["metadata"].get("visual_type", "unknown")
        stats["visual_types"][vtype] = stats["visual_types"].get(vtype, 0) + 1

    return stats