spectre_pdf 1.0.0

"""V2: add multi-line title joining to each primitive.
A heading "candidate" is a run of consecutive lines that all satisfy the
primitive's predicate; we emit both individual lines AND the joined run.
This matches what pymupdf4llm effectively does."""
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


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, gt):
    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):
    out = set()
    for entry in doc.get_toc():
        h = normalize_heading(entry[1])
        if h and len(h) > 1:
            out.add(h)
    return out


def collect_spans(doc):
    pages = []
    for page in doc:
        d = page.get_text("dict")
        spans = []
        for block in d.get("blocks", []):
            if block.get("type") != 0:
                continue
            for line in block.get("lines", []):
                for sp in line.get("spans", []):
                    if not sp.get("text", "").strip():
                        continue
                    b = sp["bbox"]
                    spans.append({
                        "x0": b[0], "y0": b[1], "x1": b[2], "y1": b[3],
                        "size": round(sp.get("size", 0), 1),
                        "flags": sp.get("flags", 0),
                        "text": sp["text"],
                    })
        pages.append(spans)
    return pages


def body_baseline(all_spans):
    c = 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, y_tol=2.0):
    if not spans:
        return []
    by_y = sorted(spans, key=lambda s: s["y0"])
    lines, cur, cur_y = [], [by_y[0]], 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, cur_y = [s], s["y0"]
    cur.sort(key=lambda x: x["x0"])
    lines.append(cur)
    return lines


def line_text(line):
    return " ".join(s["text"] for s in line).strip()


def line_size(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):
    tot = sum(len(s["text"]) for s in line)
    if tot == 0:
        return 0.0
    return sum(len(s["text"]) for s in line if s["flags"] & BOLD_FLAG) / tot


def emit_runs(lines, is_heading_flags):
    """Yield (joined_text, [line_indices]) for maximal consecutive runs."""
    runs = []
    i = 0
    n = len(lines)
    while i < n:
        if is_heading_flags[i]:
            j = i
            while j + 1 < n and is_heading_flags[j + 1]:
                j += 1
            joined = " ".join(line_text(lines[k]) for k in range(i, j + 1))
            runs.append(joined)
            i = j + 1
        else:
            i += 1
    return runs


# Primitive A: vertical-gap + size
def primitive_a(all_spans, baseline):
    out = 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))]
        pos = [g for g in gaps if g > 0]
        if not pos:
            continue
        med = median(pos)
        flags = []
        for i, line in enumerate(lines):
            gap_above = gaps[i - 1] if i > 0 else med * 3
            sz = line_size(line)
            flags.append(gap_above > 2 * med or sz > baseline + 0.5)
        for run in emit_runs(lines, flags):
            h = normalize_heading(run)
            if h and len(h) > 1 and len(run.split()) <= 30:
                out.add(h)
    return out


# Primitive B: font-flag
def primitive_b(all_spans):
    out = set()
    terminal = re.compile(r"[\.\?\!]$")
    for page_spans in all_spans:
        lines = group_into_lines(page_spans)
        flags = []
        for line in lines:
            t = line_text(line)
            ok = (line_bold_frac(line) >= 0.8
                  and 0 < len(t.split()) <= 16
                  and not terminal.search(t.rstrip()))
            flags.append(ok)
        for run in emit_runs(lines, flags):
            if len(run.split()) > 30:
                continue
            h = normalize_heading(run)
            if h and len(h) > 1:
                out.add(h)
    return out


# Primitive C: hybrid score
def primitive_c(all_spans, baseline):
    out = 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))]
        pos = [g for g in gaps if g > 0]
        if not pos:
            continue
        med = median(pos)
        flags = []
        for i, line in enumerate(lines):
            sc = 0
            if line_size(line) > baseline + 0.5:
                sc += 1
            if line_bold_frac(line) >= 0.99:
                sc += 1
            gap_above = gaps[i - 1] if i > 0 else med * 3
            if gap_above > 1.5 * med:
                sc += 1
            flags.append(sc >= 2)
        for run in emit_runs(lines, flags):
            if len(run.split()) > 30:
                continue
            h = normalize_heading(run)
            if h and len(h) > 1:
                out.add(h)
    return out


def main():
    doc = pymupdf.open(PDF)
    gt = gt_headings(doc)
    all_spans = collect_spans(doc)
    baseline = body_baseline(all_spans)
    print(f"PDF: {PDF.name}  pages={len(doc)}  TOC={len(gt)}  baseline_size={baseline}")
    print()
    results = {}
    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)
        results[name] = (len(det), p, r, f1)
        print(f"{name:<22} detections={len(det):<5} "
              f"P={p:.3f}  R={r:.3f}  F1={f1:.3f}")

    print()
    best = max(results.items(), key=lambda kv: kv[1][3])
    print(f"BEST: {best[0]}  F1={best[1][3]:.3f}")


if __name__ == "__main__":
    main()