backend/scripts/diagnose_drops.py

"""Diagnostic: trace where parse_pdf drops content, per gate.

Mirrors parse_pdf/segment_markdown exactly (importing the real helpers) but
logs every dropped line/block with the gate that killed it. Read-only.

Usage: uv run python scripts/diagnose_drops.py ../test_papers/foo.pdf
"""
import re
import sys
from collections import Counter
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from astroparse_api.parse import (  # noqa: E402
    _AFFILIATION, _AUTHOR_BLOCK_FRONT_LIMIT, _CAPTION, _HEADING, _MIN_PARA,
    _PICTURE, _REFERENCES, _TABLE_ROW, _norm, _pathological_pages,
    _is_author_affiliation_block, _prev_para_is_open, demarkdown,
    repair_hyphenation, strip_picture_text, _PAGE_NUM, _ARXIV_STAMP,
    _DOI_STAMP, _RUNNING_HEADER, _fallback_chunk,
)


def get_markdown(pdf_path: Path) -> str:
    import pymupdf
    import pymupdf4llm

    doc = pymupdf.open(str(pdf_path))
    n_pages = doc.page_count
    bad = _pathological_pages(doc)
    if bad:
        normal_pages = [i for i in range(n_pages) if i not in bad]
        md_chunks = pymupdf4llm.to_markdown(doc, page_chunks=True, pages=normal_pages, force_text=False)
        it = iter(md_chunks)
        chunks = [_fallback_chunk(doc[i]) if i in bad else next(it) for i in range(n_pages)]
    else:
        chunks = pymupdf4llm.to_markdown(doc, page_chunks=True, force_text=False)

    pages_lines = [c["text"].splitlines() for c in chunks]

    # --- remove_repeating_lines, instrumented ---
    counts = Counter()
    for page in pages_lines:
        for n in {_norm(l) for l in page if l.strip()}:
            counts[n] += 1
    threshold = max(2, len(pages_lines) // 2)
    repeating = {n for n, c in counts.items() if c >= threshold}
    removed_rep = Counter()
    kept_pages = []
    for page in pages_lines:
        kept = []
        for l in page:
            if _norm(l) in repeating:
                removed_rep[l.strip()] += 1
            else:
                kept.append(l)
        kept_pages.append(kept)
    print("=== GATE: remove_repeating_lines ===")
    for line, c in removed_rep.most_common(40):
        print(f"  [{c}x] {line[:110]}")

    # --- strip_page_artifacts, instrumented ---
    print("=== GATE: strip_page_artifacts ===")
    out_pages = []
    for page in kept_pages:
        kept = []
        for l in page:
            s = l.strip()
            if not s:
                kept.append(l)
                continue
            if _PAGE_NUM.match(s):
                pass  # page numbers: uninteresting
            elif _ARXIV_STAMP.search(l):
                print(f"  [arxiv-stamp] {s[:110]}")
            elif _DOI_STAMP.match(s):
                print(f"  [doi] {s[:110]}")
            elif _RUNNING_HEADER.match(s):
                print(f"  [RUNNING_HEADER] {s[:110]}")
            else:
                kept.append(l)
                continue
        out_pages.append(kept)

    out_pages = [strip_picture_text(p) for p in out_pages]
    return "\n".join("\n".join(p) for p in out_pages)


def segment_instrumented(md: str):
    paras, section, first_pending = [], "", True
    print("=== GATE: segment_markdown ===")
    for block in re.split(r"\n\s*\n", md):
        block = repair_hyphenation(block).strip()
        if not block:
            continue
        m = _HEADING.match(block)
        if m:
            name = m.group(1).strip().strip("*")
            rest = block[m.end():].strip()
            if _REFERENCES.match(name):
                tail = len(md) - md.find(block)
                print(f"  [BREAK at references-like heading] '{name}' — discarding remaining {tail} chars of doc")
                break
            if rest:
                print(f"  [HEADING-GLUED TEXT DROPPED] heading='{name[:60]}' dropped {len(rest)} chars: {rest[:140]!r}")
            section, first_pending = name, True
            continue
        filtered_lines = strip_picture_text([l for l in block.splitlines() if not _AFFILIATION.match(l)])
        block = "\n".join(filtered_lines).strip()
        if not block:
            continue
        text = re.sub(r"\s+", " ", block)
        if _CAPTION.match(text) or _PICTURE.match(text) or _TABLE_ROW.match(text):
            gate = 'CAPTION' if _CAPTION.match(text) else ('PICTURE' if _PICTURE.match(text) else 'TABLE')
            # captions are legitimately long; only flag caption-gate kills that
            # look like prose (verb right after the figure number)
            if gate == 'CAPTION' and re.match(r'\*{0,2}(figure|fig\.|table)s?\s*\d+[a-z\s]', text, re.I) and not re.match(r'\*{0,2}(figure|fig\.|table)s?\s*\d+\s*[.:|]', text, re.I):
                print(f"  [CAPTION-GATE PROSE? DROPPED] ({len(text)} ch) sec={section[:40]!r}: {text[:130]!r}")
            elif gate != 'CAPTION' and len(text) > 300:
                print(f"  [{gate}-GATE DROPPED] ({len(text)} ch): {text[:120]!r}")
            continue
        if len(paras) < _AUTHOR_BLOCK_FRONT_LIMIT and _is_author_affiliation_block(text):
            print(f"  [author-block] dropped {len(text)} chars: {text[:100]!r}")
            continue
        starts_lower = text[:1].islower()
        prev_open = bool(paras and _prev_para_is_open(paras[-1]["text"]))
        should_merge = starts_lower or len(text) < _MIN_PARA or prev_open
        if paras and should_merge and paras[-1]["section"] == section:
            paras[-1]["text"] += " " + demarkdown(text)
        elif len(text) >= _MIN_PARA:
            paras.append({"section": section, "text": demarkdown(text)})
            first_pending = False
        else:
            why = []
            if not paras:
                why.append("no prior para")
            elif paras[-1]["section"] != section:
                why.append(f"section changed ({paras[-1]['section'][:30]!r} -> {section[:30]!r})")
            if not should_merge:
                why.append("no merge trigger")
            print(f"  [SHORT-BLOCK DROPPED] ({len(text)} ch, {', '.join(why)}) sec={section[:40]!r}: {text[:120]!r}")
    short_final = [p for p in paras if len(p["text"]) < _MIN_PARA]
    for p in short_final:
        print(f"  [FINAL-FILTER DROPPED] ({len(p['text'])} ch) sec={p['section'][:40]!r}: {p['text'][:120]!r}")
    paras = [p for p in paras if len(p["text"]) >= _MIN_PARA]
    print(f"=== RESULT: {len(paras)} paragraphs, sections seen: ===")
    seen = []
    for p in paras:
        if p["section"] not in seen:
            seen.append(p["section"])
    for s in seen:
        print(f"  - {s[:90]!r}")
    return paras


if __name__ == "__main__":
    pdf = Path(sys.argv[1])
    print(f"##### {pdf.name} #####")
    md = get_markdown(pdf)
    # Also: what do raw headings look like before any filtering?
    print("=== RAW HEADING-LIKE LINES in markdown ===")
    for line in md.splitlines():
        s = line.strip()
        if s.startswith("#"):
            print(f"  {s[:100]}")
    segment_instrumented(md)