spectre_pdf 1.0.0

"""
Analyze segmentation differences between spectre-rs and mupdf (via pymupdf)
that cause heading-detection misses.

For each /Outline TOC entry on i1040gi:
  1. Find where the heading text lives in pymupdf's get_text("dict") output
     (which block, line, span; what font size).
  2. Find where the heading text lives in spectre_rs.Document.blocks()
     output (which block, line, span; what font size).
  3. Categorize the difference:
       - SAME: both segment the heading identically
       - COL_CUT: spectre splits heading across blocks via column pre-pass
       - LINE_SPLIT: spectre splits heading across lines
       - MERGED_WITH_BODY: spectre puts heading text inside a body block
       - MISSING: spectre doesn't have the heading text at all
       - FONT_SIZE_DIFF: spectre reads a different font size

Output: a frequency table of difference categories so we know what targeted
fixes would close the most heading-detection misses.
"""

from __future__ import annotations

import re
import sys
from collections import Counter
from pathlib import Path

import fitz
import spectre_rs


def norm(s: str) -> str:
    s = re.sub(r"[*_`]", "", s or "")
    s = re.sub(r"\s+", " ", s).strip().lower()
    return re.sub(r"[\.\:\,\;\!\?\—\-]+$", "", s)


def main() -> None:
    pdf_path = "/mnt/c/Users/RyanJ/spectre-rs/bench-corpus/largedoc/i1040gi.pdf"
    pdoc = fitz.open(pdf_path)
    sd = spectre_rs.Document(open(pdf_path, "rb").read())

    toc = [(t[0], t[1], t[2]) for t in pdoc.get_toc()]
    toc_targets: dict[str, tuple[int, int]] = {}
    for lvl, title, page in toc:
        if title.strip():
            toc_targets[norm(title)] = (lvl, page)

    # Build per-page indexes from pymupdf
    pmu_lines_by_page: dict[int, list[tuple[str, float, str]]] = {}
    for page_no in range(pdoc.page_count):
        page = pdoc[page_no]
        d = page.get_text("dict")
        lines = []
        for blk in d["blocks"]:
            if blk.get("type") != 0:
                continue
            for line in blk.get("lines", []):
                line_text = "".join(s["text"] for s in line.get("spans", []))
                if not line_text.strip():
                    continue
                max_size = max(
                    (s["size"] for s in line["spans"]), default=0
                )
                flags = (
                    line["spans"][0]["flags"] if line["spans"] else 0
                )
                font = line["spans"][0]["font"] if line["spans"] else ""
                lines.append((line_text, max_size, font))
        pmu_lines_by_page[page_no + 1] = lines

    # Build per-page indexes from spectre
    spec_blocks_by_page: dict[int, list] = {}
    for block in sd.blocks():
        spec_blocks_by_page.setdefault(block.page, []).append(block)

    categories: Counter[str] = Counter()
    examples: dict[str, list[tuple[str, str, str]]] = {
        k: [] for k in [
            "SAME", "COL_CUT", "LINE_SPLIT", "MERGED_WITH_BODY",
            "MISSING", "FONT_SIZE_DIFF",
        ]
    }

    for title_norm, (lvl, page) in toc_targets.items():
        if len(title_norm) < 4:
            continue
        # Look at the next ~3 pages since /Outline page numbers can be
        # off by 1 or interior.
        candidate_pages = [page] if page in pmu_lines_by_page else []
        for p in [page - 1, page, page + 1, page + 2]:
            if p in pmu_lines_by_page and p not in candidate_pages:
                candidate_pages.append(p)

        # Find pymupdf line(s) containing this heading
        pmu_match = None
        for p in candidate_pages:
            for line_text, size, font in pmu_lines_by_page.get(p, []):
                if title_norm in norm(line_text):
                    pmu_match = (p, line_text, size, font)
                    break
            if pmu_match:
                break

        # Find spectre block/line containing this heading
        spec_match = None
        spec_partial_blocks: list = []
        for p in candidate_pages:
            for block in spec_blocks_by_page.get(p, []):
                if title_norm in norm(block.text):
                    # Find specific line within block
                    for line in block.lines:
                        if title_norm in norm(line.text):
                            spec_match = (
                                p, line.text,
                                max((s.font_size for s in line.spans), default=0),
                                block,
                            )
                            break
                    if spec_match:
                        break
                # Partial match: heading split across blocks?
                title_words = title_norm.split()
                if title_words:
                    first_words = " ".join(title_words[:2])
                    if first_words in norm(block.text):
                        spec_partial_blocks.append((p, block))
            if spec_match:
                break

        # Categorize
        if not pmu_match:
            continue  # pmu can't find it either, skip
        if not spec_match:
            if spec_partial_blocks:
                categories["COL_CUT"] += 1
                if len(examples["COL_CUT"]) < 5:
                    examples["COL_CUT"].append((
                        title_norm,
                        pmu_match[1][:60],
                        f"partials: {len(spec_partial_blocks)} blocks",
                    ))
            else:
                categories["MISSING"] += 1
                if len(examples["MISSING"]) < 5:
                    examples["MISSING"].append((
                        title_norm,
                        pmu_match[1][:60],
                        "no match",
                    ))
            continue

        pmu_p, pmu_text, pmu_size, pmu_font = pmu_match
        spec_p, spec_text, spec_size, spec_block = spec_match
        if abs(pmu_size - spec_size) > 1.5:
            categories["FONT_SIZE_DIFF"] += 1
            if len(examples["FONT_SIZE_DIFF"]) < 5:
                examples["FONT_SIZE_DIFF"].append((
                    title_norm,
                    f"pmu {pmu_size:.0f}pt",
                    f"spec {spec_size:.0f}pt",
                ))
            continue
        # Both have it, similar size. Now check whether spec's block is
        # "pure heading" (single line) or "heading + body merged".
        spec_dom_block_size = 0.0
        if spec_block.lines:
            char_by_size: Counter[float] = Counter()
            for ln in spec_block.lines:
                for sp in ln.spans:
                    char_by_size[round(sp.font_size)] += len(sp.text)
            if char_by_size:
                spec_dom_block_size = char_by_size.most_common(1)[0][0]
        if spec_dom_block_size and abs(spec_dom_block_size - spec_size) > 1.5:
            categories["MERGED_WITH_BODY"] += 1
            if len(examples["MERGED_WITH_BODY"]) < 5:
                examples["MERGED_WITH_BODY"].append((
                    title_norm,
                    f"spec line {spec_size:.0f}pt",
                    f"block dom {spec_dom_block_size:.0f}pt",
                ))
            continue
        if len(spec_block.lines) > 1:
            categories["LINE_SPLIT"] += 1
            if len(examples["LINE_SPLIT"]) < 5:
                examples["LINE_SPLIT"].append((
                    title_norm,
                    pmu_text[:60],
                    f"spec block has {len(spec_block.lines)} lines",
                ))
            continue
        categories["SAME"] += 1
        if len(examples["SAME"]) < 5:
            examples["SAME"].append((
                title_norm, pmu_text[:60], f"{spec_size:.0f}pt",
            ))

    print("=" * 80)
    print("Segmentation diff: spectre-rs vs pymupdf on i1040gi.pdf")
    print("=" * 80)
    print(f"Total TOC entries analyzed: {sum(categories.values())}")
    print()
    for cat in [
        "SAME", "MERGED_WITH_BODY", "COL_CUT", "LINE_SPLIT",
        "FONT_SIZE_DIFF", "MISSING",
    ]:
        count = categories[cat]
        pct = count / max(sum(categories.values()), 1) * 100
        print(f"  {cat:<20} {count:>5}  ({pct:>5.1f}%)")
    print()
    print("Examples per category:")
    for cat, ex in examples.items():
        if not ex:
            continue
        print(f"\n  {cat}:")
        for title, a, b in ex:
            print(f"    {title!r:<50} | {a} | {b}")


if __name__ == "__main__":
    main()