spectre_pdf 1.0.0

"""
Throwaway: test 3 heading-segmentation primitives on i1040gi.pdf using
pymupdf raw spans as input, scored against /Outline TOC ground truth.

Primitives:
  A: vertical-gap segmentation (gap > 2x median OR size > body baseline)
  B: pure font-flag (bold + <=16 words + no terminal punctuation)
  C: hybrid score (size + bold + gap), threshold >= 2

NOT integrated with spectre-rs. Span data comes from page.get_text("dict").
"""
from __future__ import annotations

import re
from collections import Counter
from pathlib import Path
from statistics import median

import pymupdf

PDF = Path("/mnt/c/Users/RyanJ/spectre-rs/bench-corpus/largedoc/i1040gi.pdf")
BOLD_FLAG = 16  # pymupdf span flag bit 4 == bold


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


def prec_recall_f1(detected: set[str], gt: set[str]) -> tuple[float, float, float]:
    if not gt:
        return 0.0, 0.0, 0.0
    tp = len(detected & gt)
    p = tp / len(detected) if detected else 0.0
    r = tp / len(gt)
    f1 = 2 * p * r / (p + r) if (p + r) else 0.0
    return p, r, f1


def gt_headings(doc: pymupdf.Document) -> set[str]:
    out: set[str] = set()
    for entry in doc.get_toc():
        title = entry[1]
        h = normalize_heading(title)
        if h and len(h) > 1:
            out.add(h)
    return out


def collect_spans(doc: pymupdf.Document) -> list[list[dict]]:
    """Return [page_idx -> list of span dicts with x0,y0,x1,y1,size,flags,text,font]."""
    pages = []
    for page in doc:
        d = page.get_text("dict")
        spans = []
        for block in d.get("blocks", []):
            if block.get("type") != 0:  # text blocks only
                continue
            for line in block.get("lines", []):
                for sp in line.get("spans", []):
                    txt = sp.get("text", "")
                    if not txt.strip():
                        continue
                    bbox = sp.get("bbox", (0, 0, 0, 0))
                    spans.append({
                        "x0": bbox[0], "y0": bbox[1],
                        "x1": bbox[2], "y1": bbox[3],
                        "size": round(sp.get("size", 0), 1),
                        "flags": sp.get("flags", 0),
                        "font": sp.get("font", ""),
                        "text": txt,
                    })
        pages.append(spans)
    return pages


def body_baseline(all_spans: list[list[dict]]) -> float:
    """Char-weighted dominant font size."""
    c: Counter[float] = Counter()
    for page in all_spans:
        for s in page:
            c[s["size"]] += len(s["text"])
    return c.most_common(1)[0][0] if c else 10.0


def group_into_lines(spans: list[dict], y_tol: float = 2.0) -> list[list[dict]]:
    """Group spans on same baseline; return list of lines sorted top-to-bottom."""
    if not spans:
        return []
    by_y = sorted(spans, key=lambda s: s["y0"])
    lines: list[list[dict]] = []
    cur: list[dict] = [by_y[0]]
    cur_y = by_y[0]["y0"]
    for s in by_y[1:]:
        if abs(s["y0"] - cur_y) <= y_tol:
            cur.append(s)
        else:
            cur.sort(key=lambda x: x["x0"])
            lines.append(cur)
            cur = [s]
            cur_y = s["y0"]
    cur.sort(key=lambda x: x["x0"])
    lines.append(cur)
    return lines


def line_text(line: list[dict]) -> str:
    return " ".join(s["text"] for s in line).strip()


def line_size(line: list[dict]) -> float:
    """Char-weighted size of a line."""
    tot = sum(len(s["text"]) for s in line)
    if tot == 0:
        return 0.0
    return sum(s["size"] * len(s["text"]) for s in line) / tot


def line_bold_frac(line: list[dict]) -> float:
    tot = sum(len(s["text"]) for s in line)
    if tot == 0:
        return 0.0
    bold = sum(len(s["text"]) for s in line if s["flags"] & BOLD_FLAG)
    return bold / tot


# ---------------- Primitive A: vertical-gap ----------------
def primitive_a(all_spans: list[list[dict]], baseline: float) -> set[str]:
    out: set[str] = set()
    for page_spans in all_spans:
        lines = group_into_lines(page_spans)
        if len(lines) < 2:
            continue
        gaps = [lines[i][0]["y0"] - lines[i - 1][0]["y0"] for i in range(1, len(lines))]
        positive = [g for g in gaps if g > 0]
        if not positive:
            continue
        med = median(positive)
        for i, line in enumerate(lines):
            gap_above = gaps[i - 1] if i > 0 else med * 3
            sz = line_size(line)
            is_heading = (gap_above > 2 * med) or (sz > baseline + 0.5)
            if not is_heading:
                continue
            t = line_text(line)
            h = normalize_heading(t)
            if h and len(h) > 1 and len(t.split()) <= 25:
                out.add(h)
    return out


# ---------------- Primitive B: font-flag only ----------------
def primitive_b(all_spans: list[list[dict]]) -> set[str]:
    out: set[str] = set()
    terminal = re.compile(r"[\.\?\!]$")
    for page_spans in all_spans:
        for line in group_into_lines(page_spans):
            t = line_text(line)
            if not t:
                continue
            if line_bold_frac(line) < 0.8:
                continue
            words = t.split()
            if len(words) > 16 or len(words) < 1:
                continue
            if terminal.search(t.rstrip()):
                continue
            h = normalize_heading(t)
            if h and len(h) > 1:
                out.add(h)
    return out


# ---------------- Primitive C: hybrid score ----------------
def primitive_c(all_spans: list[list[dict]], baseline: float) -> set[str]:
    out: set[str] = set()
    for page_spans in all_spans:
        lines = group_into_lines(page_spans)
        if len(lines) < 2:
            continue
        gaps = [lines[i][0]["y0"] - lines[i - 1][0]["y0"] for i in range(1, len(lines))]
        positive = [g for g in gaps if g > 0]
        if not positive:
            continue
        med = median(positive)
        for i, line in enumerate(lines):
            score = 0
            if line_size(line) > baseline + 0.5:
                score += 1
            if line_bold_frac(line) >= 0.99:
                score += 1
            gap_above = gaps[i - 1] if i > 0 else med * 3
            if gap_above > 1.5 * med:
                score += 1
            if score < 2:
                continue
            t = line_text(line)
            if not t or len(t.split()) > 25:
                continue
            h = normalize_heading(t)
            if h and len(h) > 1:
                out.add(h)
    return out


def main() -> None:
    doc = pymupdf.open(PDF)
    gt = gt_headings(doc)
    print(f"PDF: {PDF.name}  pages={len(doc)}  TOC entries (normalized)={len(gt)}")
    all_spans = collect_spans(doc)
    baseline = body_baseline(all_spans)
    print(f"Body baseline font size (char-weighted dominant): {baseline}")
    print()

    for name, fn in [
        ("A: vertical-gap", lambda: primitive_a(all_spans, baseline)),
        ("B: font-flag",    lambda: primitive_b(all_spans)),
        ("C: hybrid score", lambda: primitive_c(all_spans, baseline)),
    ]:
        det = fn()
        p, r, f1 = prec_recall_f1(det, gt)
        print(f"{name:<22} detections={len(det):<5} "
              f"P={p:.3f}  R={r:.3f}  F1={f1:.3f}")


if __name__ == "__main__":
    main()