refactor: add local reference zones and layout-aware tail boundary detection

paperforge/worker/ocr_document.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import dataclasses
 from collections import namedtuple
 from dataclasses import dataclass, field
 
@@ -57,6 +58,15 @@ class ReadingSegment:
     semantic_hint: str = ""
 
 
+@dataclass
+class ReferenceZone:
+    page: int
+    column_index: int
+    y_start: float
+    y_end: float
+    block_indices: list[int]
+
+
 @dataclass
 class DocumentStructure:
     body_end_page: int | None = None
@@ -67,6 +77,7 @@ class DocumentStructure:
     backmatter_form: str = "flat"
     page_layouts: dict[int, PageLayoutProfile] | None = None
     tail_reading_order: list[dict] | None = None
+    reference_zones: list[dict] | None = None
 
 
 def _cluster_page_columns(page_blocks: list[dict], page_width: float) -> list[float]:
@@ -310,13 +321,86 @@ def _build_tail_reading_order(
     return segments
 
 
-def _detect_forward_body_end(blocks: list[dict]) -> int | None:
+def _detect_reference_zones(
+    blocks: list[dict],
+    page_layouts: dict[int, PageLayoutProfile],
+) -> list[ReferenceZone]:
+    """Detect local reference zones per page.
+
+    For each page with a ``reference_heading``, creates a ``ReferenceZone``
+    scoped to that heading's column.  Only blocks in the same column and
+    below the heading's bottom y are included in the zone.
+    """
+    zones: list[ReferenceZone] = []
+    for block in blocks:
+        if block.get("role") != "reference_heading":
+            continue
+        page = block.get("page")
+        if page is None:
+            continue
+        profile = page_layouts.get(page)
+        bbox = _block_bbox(block)
+        if not bbox:
+            continue
+        y_start = bbox[3]
+        if profile is None or profile.column_count <= 1:
+            column_index = 0
+        else:
+            x_center = (bbox[0] + bbox[2]) / 2
+            column_index = _get_column_index_by_boundaries(x_center, profile.column_boundaries)
+        zone_indices: list[int] = []
+        max_y = y_start
+        for j, b in enumerate(blocks):
+            if b.get("page") != page:
+                continue
+            bb = _block_bbox(b)
+            if not bb:
+                continue
+            if profile and profile.column_count > 1:
+                bx_center = (bb[0] + bb[2]) / 2
+                bcol = _get_column_index_by_boundaries(bx_center, profile.column_boundaries)
+                if bcol != column_index:
+                    continue
+            if bb[1] < y_start - 1:
+                continue
+            zone_indices.append(j)
+            if bb[3] > max_y:
+                max_y = bb[3]
+        zones.append(
+            ReferenceZone(
+                page=page,
+                column_index=column_index,
+                y_start=y_start,
+                y_end=max_y,
+                block_indices=zone_indices,
+            )
+        )
+    return zones
+
+
+def _block_in_any_reference_zone(
+    block: dict,
+    zones: list[ReferenceZone],
+    global_index: int,
+) -> bool:
+    """Check if a block (by global index) falls within any reference zone."""
+    return any(global_index in zone.block_indices for zone in zones)
+
+
+def _detect_forward_body_end(
+    blocks: list[dict],
+    page_layouts: dict[int, PageLayoutProfile] | None = None,
+) -> int | None:
     """Scan blocks front-to-back and return the last page of stable body.
 
     Tracks pages with body headings (section_heading, subsection_heading)
     and body_paragraph continuity.  When a page has tail roles
     (backmatter_heading, reference_heading, etc.) and no body content,
     the body is considered to have ended on the preceding clean body page.
+
+    With ``page_layouts``, multi-column pages are checked per-column:
+    if any column has body roles without tail roles in that same column,
+    body continues even if other columns have tail content.
     Returns None if no clear body/backmatter boundary is found.
     """
     if not blocks:
@@ -330,13 +414,41 @@ def _detect_forward_body_end(blocks: list[dict]) -> int | None:
     if not pages:
         return None
 
+    body_roles = {"body_paragraph", "section_heading", "subsection_heading", "sub_subsection_heading"}
     last_clean_body_page: int | None = None
 
     for page in pages:
-        roles = {b.get("role") for b in by_page[page]}
-        has_body = bool(roles & {"body_paragraph", "section_heading", "subsection_heading", "sub_subsection_heading"})
+        page_blocks = by_page[page]
+        roles = {b.get("role") for b in page_blocks}
+        has_body = bool(roles & body_roles)
         has_tail = bool(roles & _TAIL_ROLES)
 
+        if page_layouts and page in page_layouts:
+            profile = page_layouts[page]
+            if profile.column_count > 1:
+                boundaries = profile.column_boundaries
+                col_has_body: dict[int, bool] = {}
+                col_has_tail: dict[int, bool] = {}
+                for block in page_blocks:
+                    bbox = _block_bbox(block)
+                    if not bbox:
+                        continue
+                    x_center = (bbox[0] + bbox[2]) / 2
+                    col = _get_column_index_by_boundaries(x_center, boundaries)
+                    role = block.get("role", "")
+                    if role in body_roles:
+                        col_has_body[col] = True
+                    if role in _TAIL_ROLES:
+                        col_has_tail[col] = True
+
+                any_body_without_tail = any(
+                    col_has_body.get(col, False) and not col_has_tail.get(col, False)
+                    for col in range(profile.column_count)
+                )
+                if any_body_without_tail:
+                    last_clean_body_page = page
+                    continue
+
         if has_body and not has_tail:
             last_clean_body_page = page
         elif has_tail:
@@ -348,16 +460,23 @@ def _detect_forward_body_end(blocks: list[dict]) -> int | None:
                     return pages[prev_idx]
                 return None
 
-    return None
+    return last_clean_body_page
 
 
-def _detect_backward_backmatter_start(blocks: list[dict]) -> int | None:
+def _detect_backward_backmatter_start(
+    blocks: list[dict],
+    page_layouts: dict[int, PageLayoutProfile] | None = None,
+) -> int | None:
     """Scan blocks backward and return the page where backmatter begins.
 
     Starting from the last page, looks for the first reference_heading or
     backmatter_heading.  Dense reference pages (>= 4 reference_item blocks)
     are a strong signal.  Short backmatter_body blocks near headings confirm
-    the backmatter zone.  Returns None if no backmatter found.
+    the backmatter zone.
+
+    With ``page_layouts``, multi-column pages check for backmatter headings
+    rather than relying on the dense-refs heuristic, which can be confused
+    by reference-only columns.  Returns None if no backmatter found.
     """
     if not blocks:
         return None
@@ -377,6 +496,11 @@ def _detect_backward_backmatter_start(blocks: list[dict]) -> int | None:
         if "reference_heading" in roles or "backmatter_heading" in roles or "backmatter_boundary_heading" in roles:
             return page
 
+        if page_layouts and page in page_layouts:
+            profile = page_layouts[page]
+            if profile.column_count > 1:
+                continue
+
         dense_refs = sum(1 for b in page_blocks if b.get("role") == "reference_item")
         if dense_refs >= 4:
             return page
@@ -384,11 +508,26 @@ def _detect_backward_backmatter_start(blocks: list[dict]) -> int | None:
     return None
 
 
-def _detect_references_start(blocks: list[dict], body_end_page: int | None) -> int | None:
+def _detect_references_start(
+    blocks: list[dict],
+    body_end_page: int | None,
+    page_layouts: dict[int, PageLayoutProfile] | None = None,
+) -> int | None:
     """Scan from body end page forward for the first page with a reference
-    heading or reference item.  Returns None if no references zone is found."""
+    heading or reference item.
+
+    With ``page_layouts``, uses ``_detect_reference_zones`` internally to
+    find the earliest page with a local reference zone.  Falls back to
+    page-level scanning when layout data is unavailable.
+    Returns None if no references zone is found.
+    """
     if body_end_page is None:
         return None
+    if page_layouts:
+        zones = _detect_reference_zones(blocks, page_layouts)
+        if zones:
+            earliest = min(z.page for z in zones)
+            return max(earliest, body_end_page)
     by_page: dict[int, list[dict]] = {}
     for block in blocks:
         p = block.get("page")
@@ -402,15 +541,19 @@ def _detect_references_start(blocks: list[dict], body_end_page: int | None) -> i
     return None
 
 
-def _reconcile_tail_spread(blocks: list[dict]) -> TailBoundary | None:
+def _reconcile_tail_spread(
+    blocks: list[dict],
+    page_layouts: dict[int, PageLayoutProfile] | None = None,
+) -> TailBoundary | None:
     """Reconcile forward and backward scans into a structured TailBoundary.
 
     Returns a TailBoundary namedtuple or None when no tail spread exists.
     The ``reason`` field provides an explainability trace.
+    ``page_layouts`` is passed to layout-aware boundary detection functions.
     """
-    forward_end = _detect_forward_body_end(blocks)
-    backward_start = _detect_backward_backmatter_start(blocks)
-    references_start = _detect_references_start(blocks, forward_end)
+    forward_end = _detect_forward_body_end(blocks, page_layouts)
+    backward_start = _detect_backward_backmatter_start(blocks, page_layouts)
+    references_start = _detect_references_start(blocks, forward_end, page_layouts)
 
     if forward_end is None and backward_start is None:
         return None
@@ -892,7 +1035,8 @@ def analyze_document_structure(blocks: list[dict]) -> DocumentStructure:
     Uses the existing _detect_forward_body_end, _detect_backward_backmatter_start,
     etc. internally. Returns a DocumentStructure with all boundary info.
     """
-    tail_spread = _reconcile_tail_spread(blocks)
+    page_layouts = _build_page_layout_profiles(blocks)
+    tail_spread = _reconcile_tail_spread(blocks, page_layouts)
     if tail_spread is not None:
         backmatter_form = _classify_backmatter_form(tail_spread, blocks)
         _label_backmatter_regime(tail_spread, backmatter_form, blocks)
@@ -900,7 +1044,7 @@ def analyze_document_structure(blocks: list[dict]) -> DocumentStructure:
     else:
         backmatter_form = "flat"
 
-    return DocumentStructure(
+    ds = DocumentStructure(
         body_end_page=tail_spread.body_end_page if tail_spread else None,
         backmatter_start=PagePosition(page=tail_spread.backmatter_start, y=0.0)
         if tail_spread and tail_spread.backmatter_start is not None
@@ -911,8 +1055,19 @@ def analyze_document_structure(blocks: list[dict]) -> DocumentStructure:
         spread_start=tail_spread.spread_start if tail_spread else None,
         spread_end=tail_spread.spread_end if tail_spread else None,
         backmatter_form=backmatter_form,
+        page_layouts=page_layouts,
     )
 
+    reading_segments = _build_tail_reading_order(blocks, page_layouts)
+    ds.tail_reading_order = (
+        [dataclasses.asdict(seg) for seg in reading_segments] if reading_segments else None
+    )
+
+    ref_zones = _detect_reference_zones(blocks, page_layouts)
+    ds.reference_zones = [dataclasses.asdict(z) for z in ref_zones] if ref_zones else None
+
+    return ds
+
 
 def _get_column(block: dict, page_width: float = 1200) -> int:
     bbox = block.get("bbox") or block.get("block_bbox")
@@ -1423,7 +1578,9 @@ def normalize_document_structure(blocks: list[dict]) -> tuple[DocumentStructure,
     - tail body candidate promotion
     - tail spread ownership assignment
     """
-    tail_spread = _reconcile_tail_spread(blocks)
+    page_layouts = _build_page_layout_profiles(blocks)
+
+    tail_spread = _reconcile_tail_spread(blocks, page_layouts)
     if tail_spread is not None:
         backmatter_form = _classify_backmatter_form(tail_spread, blocks)
         _label_backmatter_regime(tail_spread, backmatter_form, blocks)
@@ -1442,18 +1599,19 @@ def normalize_document_structure(blocks: list[dict]) -> tuple[DocumentStructure,
         spread_start=tail_spread.spread_start if tail_spread else None,
         spread_end=tail_spread.spread_end if tail_spread else None,
         backmatter_form=backmatter_form,
+        page_layouts=page_layouts,
     )
 
     header_band, footer_band = _estimate_noise_bands(blocks)
 
-    page_layouts = _build_page_layout_profiles(blocks)
-    doc_structure.page_layouts = page_layouts
-
     reading_segments = _build_tail_reading_order(blocks, page_layouts)
     doc_structure.tail_reading_order = (
         [dataclasses.asdict(seg) for seg in reading_segments] if reading_segments else None
     )
 
+    ref_zones = _detect_reference_zones(blocks, page_layouts)
+    doc_structure.reference_zones = [dataclasses.asdict(z) for z in ref_zones] if ref_zones else None
+
     blocks = _promote_tail_body_candidates(blocks, doc_structure, header_band=header_band, footer_band=footer_band)
     blocks = _assign_tail_spread_ownership(blocks, doc_structure)
 

tests/test_ocr_document.py

@@ -1212,6 +1212,185 @@ def test_tail_reading_order_mixed_page() -> None:
     assert segments[1].semantic_hint == "mixed"
 
 
+# ---------------------------------------------------------------------------
+# Reference zone tests (Task 3)
+# ---------------------------------------------------------------------------
+
+
+def test_reference_zone_single_column() -> None:
+    """Single-column page: zone covers all blocks below ref heading."""
+    from paperforge.worker.ocr_document import (
+        _build_page_layout_profiles,
+        _detect_reference_zones,
+    )
+
+    blocks = [
+        {"page": 1, "bbox": [100, 100, 700, 140], "role": "body_paragraph", "page_width": 800, "page_height": 1000},
+        {"page": 1, "bbox": [100, 160, 700, 200], "role": "body_paragraph", "page_width": 800, "page_height": 1000},
+        {"page": 2, "bbox": [100, 100, 700, 140], "role": "reference_heading", "page_width": 800, "page_height": 1000},
+        {"page": 2, "bbox": [100, 160, 700, 200], "role": "reference_item", "page_width": 800, "page_height": 1000},
+        {"page": 2, "bbox": [100, 220, 700, 260], "role": "reference_item", "page_width": 800, "page_height": 1000},
+    ]
+
+    page_layouts = _build_page_layout_profiles(blocks)
+    zones = _detect_reference_zones(blocks, page_layouts)
+
+    assert len(zones) == 1
+    zone = zones[0]
+    assert zone.page == 2
+    assert zone.column_index == 0
+    # Zone starts at heading bottom y → heading itself excluded
+    assert len(zone.block_indices) == 2
+    for idx in zone.block_indices:
+        assert blocks[idx].get("role") == "reference_item"
+    assert zone.y_end > zone.y_start
+
+
+def test_reference_zone_two_column_mixed() -> None:
+    """Left body + right references: zone only in right column."""
+    from paperforge.worker.ocr_document import (
+        _build_page_layout_profiles,
+        _detect_reference_zones,
+    )
+
+    blocks = [
+        {"page": 1, "bbox": [100, 100, 700, 140], "role": "body_paragraph", "page_width": 800, "page_height": 1000},
+        {"page": 2, "bbox": [50, 100, 380, 140], "role": "body_paragraph", "page_width": 800, "page_height": 1000},
+        {"page": 2, "bbox": [50, 160, 380, 200], "role": "body_paragraph", "page_width": 800, "page_height": 1000},
+        {"page": 2, "bbox": [420, 100, 750, 140], "role": "reference_heading", "page_width": 800, "page_height": 1000},
+        {"page": 2, "bbox": [420, 160, 750, 200], "role": "reference_item", "page_width": 800, "page_height": 1000},
+        {"page": 2, "bbox": [420, 220, 750, 260], "role": "reference_item", "page_width": 800, "page_height": 1000},
+    ]
+
+    page_layouts = _build_page_layout_profiles(blocks)
+    zones = _detect_reference_zones(blocks, page_layouts)
+
+    assert len(zones) == 1
+    zone = zones[0]
+    assert zone.page == 2
+    assert zone.column_index == 1
+    # Zone starts at heading bottom y → heading excluded, only ref items
+    assert len(zone.block_indices) == 2
+    for idx in zone.block_indices:
+        assert blocks[idx].get("role") == "reference_item"
+    # Left-col body blocks should NOT be in zone
+    assert all(idx not in zone.block_indices for idx in [1, 2])
+
+
+def test_block_in_reference_zone() -> None:
+    """Verify block_in_any_reference_zone correctly includes/excludes."""
+    from paperforge.worker.ocr_document import (
+        ReferenceZone,
+        _block_in_any_reference_zone,
+    )
+
+    zones = [
+        ReferenceZone(page=2, column_index=1, y_start=150, y_end=300, block_indices=[3, 4, 5]),
+    ]
+
+    assert _block_in_any_reference_zone({}, zones, 3) is True
+    assert _block_in_any_reference_zone({}, zones, 5) is True
+    assert _block_in_any_reference_zone({}, zones, 0) is False
+    assert _block_in_any_reference_zone({}, zones, 2) is False
+    assert _block_in_any_reference_zone({}, [], 0) is False
+
+
+# ---------------------------------------------------------------------------
+# Layout-aware boundary detection tests (Task 4)
+# ---------------------------------------------------------------------------
+
+
+def test_forward_body_end_mixed_page() -> None:
+    """Two-column page with left body + right tail: body should continue."""
+    from paperforge.worker.ocr_document import (
+        _build_page_layout_profiles,
+        _detect_forward_body_end,
+    )
+
+    blocks = [
+        {"page": 1, "bbox": [100, 100, 700, 140], "role": "body_paragraph", "page_width": 800, "page_height": 1000},
+        {"page": 2, "bbox": [100, 100, 700, 140], "role": "body_paragraph", "page_width": 800, "page_height": 1000},
+        {"page": 3, "bbox": [100, 100, 700, 140], "role": "body_paragraph", "page_width": 800, "page_height": 1000},
+        # Page 4: left body (clean column), right tail
+        {"page": 4, "bbox": [50, 100, 380, 140], "role": "body_paragraph", "page_width": 800, "page_height": 1000},
+        {"page": 4, "bbox": [420, 100, 750, 140], "role": "backmatter_heading", "page_width": 800, "page_height": 1000},
+        {"page": 4, "bbox": [420, 160, 750, 200], "role": "backmatter_body", "page_width": 800, "page_height": 1000},
+        # Page 5: tail only
+        {"page": 5, "bbox": [420, 100, 750, 140], "role": "reference_heading", "page_width": 800, "page_height": 1000},
+        {"page": 5, "bbox": [420, 160, 750, 200], "role": "reference_item", "page_width": 800, "page_height": 1000},
+    ]
+
+    page_layouts = _build_page_layout_profiles(blocks)
+    body_end = _detect_forward_body_end(blocks, page_layouts)
+
+    assert body_end == 4, f"Expected body_end=4 (left body continues on page 4), got {body_end}"
+
+
+def test_backward_backmatter_start_mixed_page() -> None:
+    """Two-column page: dense refs in one column should NOT be backmatter_start."""
+    from paperforge.worker.ocr_document import (
+        _build_page_layout_profiles,
+        _detect_backward_backmatter_start,
+    )
+
+    blocks = [
+        {"page": 1, "bbox": [100, 100, 700, 140], "role": "body_paragraph", "page_width": 800, "page_height": 1000},
+        # Page 5: two-column with dense refs confined to right column
+        {"page": 5, "bbox": [50, 100, 380, 140], "role": "backmatter_body", "page_width": 800, "page_height": 1000},
+        {"page": 5, "bbox": [420, 50, 750, 90], "role": "reference_item", "page_width": 800, "page_height": 1000},
+        {"page": 5, "bbox": [420, 110, 750, 150], "role": "reference_item", "page_width": 800, "page_height": 1000},
+        {"page": 5, "bbox": [420, 170, 750, 210], "role": "reference_item", "page_width": 800, "page_height": 1000},
+        {"page": 5, "bbox": [420, 230, 750, 270], "role": "reference_item", "page_width": 800, "page_height": 1000},
+        {"page": 5, "bbox": [420, 290, 750, 330], "role": "reference_item", "page_width": 800, "page_height": 1000},
+        # Page 4: left backmatter heading, right reference heading
+        {"page": 4, "bbox": [50, 100, 380, 140], "role": "backmatter_heading", "page_width": 800, "page_height": 1000},
+        {"page": 4, "bbox": [420, 100, 750, 140], "role": "reference_heading", "page_width": 800, "page_height": 1000},
+        {"page": 4, "bbox": [420, 160, 750, 200], "role": "reference_item", "page_width": 800, "page_height": 1000},
+    ]
+
+    page_layouts = _build_page_layout_profiles(blocks)
+    bm_start = _detect_backward_backmatter_start(blocks, page_layouts)
+
+    assert bm_start == 4, f"Expected backmatter_start=4 (backmatter heading on page 4), got {bm_start}"
+
+
+def test_references_start_local_zone() -> None:
+    """references_start picks page 71 but zone scopes it to right column only."""
+    from paperforge.worker.ocr_document import (
+        _build_page_layout_profiles,
+        _detect_reference_zones,
+        _detect_references_start,
+    )
+
+    blocks = [
+        {"page": 1, "bbox": [100, 100, 700, 140], "role": "body_paragraph", "page_width": 800, "page_height": 1000},
+        {"page": 2, "bbox": [100, 100, 700, 140], "role": "body_paragraph", "page_width": 800, "page_height": 1000},
+        # Page 3: left body + right references (SAN9AYVR-like)
+        {"page": 3, "bbox": [50, 100, 380, 140], "role": "body_paragraph", "page_width": 800, "page_height": 1000},
+        {"page": 3, "bbox": [50, 200, 380, 240], "role": "body_paragraph", "page_width": 800, "page_height": 1000},
+        {"page": 3, "bbox": [420, 100, 750, 140], "role": "reference_heading", "page_width": 800, "page_height": 1000},
+        {"page": 3, "bbox": [420, 160, 750, 200], "role": "reference_item", "page_width": 800, "page_height": 1000},
+        {"page": 3, "bbox": [420, 220, 750, 260], "role": "reference_item", "page_width": 800, "page_height": 1000},
+    ]
+
+    page_layouts = _build_page_layout_profiles(blocks)
+    refs_start = _detect_references_start(blocks, body_end_page=2, page_layouts=page_layouts)
+    assert refs_start == 3, f"Expected references_start=3, got {refs_start}"
+
+    zones = _detect_reference_zones(blocks, page_layouts)
+    assert len(zones) == 1
+    zone = zones[0]
+    assert zone.column_index == 1
+    # Left-col body blocks NOT in zone
+    assert 2 not in zone.block_indices
+    assert 3 not in zone.block_indices
+    # Heading starts zone but is excluded (zone starts at heading bottom y)
+    assert 4 not in zone.block_indices
+    # Ref items below heading ARE in zone
+    assert 5 in zone.block_indices
+    assert 6 in zone.block_indices
+
+
 def test_document_structure_json_serialization() -> None:
     """Verify that DocumentStructure with layout profiles serializes to valid JSON."""
     from paperforge.worker.ocr_document import DocumentStructure, PageLayoutProfile