#!/usr/bin/env python3
"""
test_docling_ocr.py — Test OCR on PeerRead ACL-2017 test PDFs.

Chạy từ thư mục docs/:
    python scripts/test_docling_ocr.py --provider docling-granite        # Granite VLM qua Kaggle (khuyến nghị)
    python scripts/test_docling_ocr.py --provider docling-granite-local  # Granite VLM chạy local (cần GPU)
    python scripts/test_docling_ocr.py --provider local                  # PyMuPDF baseline
    python scripts/test_docling_ocr.py --compare

Output được linearize về đúng template plain-text của ground truth (test/parsed_pdfs):
bảng làm phẳng (cell cách space, 1 dòng/row), công thức về unicode thường (không LaTeX),
chỉ giữ heading "# "/"## " để tách section.

HuggingFace token (tùy chọn, giúp download model ổn định hơn):
    Đặt HF_TOKEN=hf_... trong file .env ở root project, hoặc:
    set HF_TOKEN=hf_...   (Windows)
    export HF_TOKEN=hf_... (Linux/Mac)
"""

from __future__ import annotations

# Suppress TensorFlow/oneDNN noise — must be set before any import
import os
os.environ.setdefault("TF_ENABLE_ONEDNN_OPTS", "0")
os.environ.setdefault("TF_CPP_MIN_LOG_LEVEL", "3")

import argparse
import asyncio
import json
import re
import sys
import unicodedata
import warnings
warnings.filterwarnings("ignore", category=UserWarning)
warnings.filterwarnings("ignore", message=".*NumPy.*")
from difflib import SequenceMatcher
from pathlib import Path

# Load .env từ root project (chứa HF_TOKEN, ANTHROPIC_API_KEY, ...)
_root = Path(__file__).resolve().parent.parent.parent
_env_file = _root / ".env"
if _env_file.exists():
    for _line in _env_file.read_text(encoding="utf-8").splitlines():
        _line = _line.strip()
        if _line and not _line.startswith("#") and "=" in _line:
            _k, _, _v = _line.partition("=")
            os.environ.setdefault(_k.strip(), _v.strip())

sys.path.insert(0, str(_root))

from backend.pipeline import pdf2md as pdf2md_backend


# ── DoclingDocument → ground-truth-style plain text ───────────────────────────
# The ground truth in test/parsed_pdfs (GROBID/CRF) is PLAIN TEXT:
#   • tables are flattened into the section text, cells joined by spaces, one row
#     per line (NO markdown pipe-tables);
#   • formulas are plain unicode (e.g. "ri = σ(Wrxi + Urhi−1 + br)"), NOT LaTeX
#     and NOT Docling's "<!-- formula-not-decoded -->" placeholder;
#   • pictures / page headers / footers are dropped.
# `result.document.export_to_markdown()` instead emits pipe-tables, "$$...$$"
# LaTeX and image placeholders, so it does not line up with the ground truth.
# We therefore linearize the DoclingDocument ourselves to the same template,
# keeping only "# "/"## " headings so md_to_peerread can still split sections.

_GREEK = {
    r"\alpha": "α", r"\beta": "β", r"\gamma": "γ", r"\delta": "δ",
    r"\epsilon": "ε", r"\varepsilon": "ε", r"\zeta": "ζ", r"\eta": "η",
    r"\theta": "θ", r"\iota": "ι", r"\kappa": "κ", r"\lambda": "λ",
    r"\mu": "μ", r"\nu": "ν", r"\xi": "ξ", r"\pi": "π", r"\rho": "ρ",
    r"\sigma": "σ", r"\tau": "τ", r"\upsilon": "υ", r"\phi": "φ",
    r"\varphi": "φ", r"\chi": "χ", r"\psi": "ψ", r"\omega": "ω",
    r"\Gamma": "Γ", r"\Delta": "Δ", r"\Theta": "Θ", r"\Lambda": "Λ",
    r"\Sigma": "Σ", r"\Phi": "Φ", r"\Psi": "Ψ", r"\Omega": "Ω",
}
_SYMBOLS = {
    r"\times": "×", r"\cdot": "·", r"\cdots": "···", r"\ldots": "...",
    r"\leq": "≤", r"\geq": "≥", r"\neq": "≠", r"\approx": "≈",
    r"\rightarrow": "→", r"\to": "→", r"\leftarrow": "←", r"\Rightarrow": "⇒",
    r"\infty": "∞", r"\sum": "∑", r"\prod": "∏", r"\int": "∫",
    r"\partial": "∂", r"\nabla": "∇", r"\in": "∈", r"\notin": "∉",
    r"\forall": "∀", r"\exists": "∃", r"\pm": "±", r"\mp": "∓",
}
_LATEX_MAP = {**_GREEK, **_SYMBOLS}


def _latex_to_plain(latex: str) -> str:
    """Best-effort LaTeX → plain unicode, to match GROBID's embedded-text style.

    e.g.  r_i = \\sigma(W_r x_i + U_r h_{i-1} + b_r)  ->  ri = σ(Wr xi + Ur hi-1 + br)
    Not a full LaTeX renderer — just enough to make the OCR output comparable
    with the plain-text ground truth.
    """
    s = (latex or "").strip()
    if not s:
        return ""
    s = s.replace("$$", "").replace("$", "")
    s = re.sub(r"\\[,;:!>\s]", " ", s)            # spacing macros: \, \; \!
    s = s.replace(r"\left", "").replace(r"\right", "")
    for k in sorted(_LATEX_MAP, key=len, reverse=True):  # longest first
        s = s.replace(k, _LATEX_MAP[k])
    s = re.sub(r"\\frac\s*\{([^{}]*)\}\s*\{([^{}]*)\}", r"(\1)/(\2)", s)
    s = re.sub(r"\\sqrt\s*\{([^{}]*)\}", r"√(\1)", s)
    s = re.sub(
        r"\\(?:mathrm|mathbf|mathbb|mathcal|text|operatorname|textbf|textit|boldsymbol)\s*\{([^{}]*)\}",
        r"\1", s,
    )
    # subscripts / superscripts: drop the _ / ^ marker, keep the content inline
    s = re.sub(r"[_^]\{([^{}]*)\}", r"\1", s)
    s = re.sub(r"[_^]\s*(\w)", r"\1", s)
    s = re.sub(r"\\([A-Za-z]+)", r"\1", s)        # leftover commands: keep name
    s = s.replace("{", "").replace("}", "")
    return re.sub(r"\s+", " ", s).strip()


def _table_to_plain(table_item) -> str:
    """Flatten a Docling table to GROBID style: cells space-joined, one row/line."""
    try:
        df = table_item.export_to_dataframe()
    except Exception:
        return ""
    if df is None or getattr(df, "empty", True):
        return ""
    import pandas as pd

    # GROBID flattens the WHOLE table onto one line: header cells then every row,
    # all space-joined (e.g. "Corpus # words # chunks # sentences Train ... Test ...").
    cells = [str(c).strip() for c in df.columns]
    for _, row in df.iterrows():
        cells.extend("" if pd.isna(v) else str(v).strip() for v in row.tolist())
    return re.sub(r"\s+", " ", " ".join(cells)).strip()


def _docling_doc_to_text(doc) -> str:
    """Linearize a DoclingDocument to GT-style plain text (with #/## headings)."""
    from docling_core.types.doc import DocItemLabel, TableItem

    # Keep CAPTION — GROBID keeps "Table 1: ...", "Figure 2: ..." inline in the text.
    drop = {
        DocItemLabel.PAGE_HEADER, DocItemLabel.PAGE_FOOTER, DocItemLabel.PICTURE,
    }
    parts: list[str] = []
    for item, _level in doc.iterate_items():
        if isinstance(item, TableItem):
            tbl = _table_to_plain(item)
            if tbl:
                parts.append(tbl)
            continue
        label = getattr(item, "label", None)
        if label in drop:
            continue
        text = (getattr(item, "text", "") or "").strip()
        # Never leak HTML comments such as the VLM's "<!-- formula-not-decoded -->"
        # placeholder — strip them so undecoded formulas simply vanish.
        text = re.sub(r"<!--.*?-->", "", text, flags=re.S).strip()
        if label == DocItemLabel.FORMULA:
            text = _latex_to_plain(text)
            if text:
                parts.append(text)
            continue
        if not text:
            continue
        if label == DocItemLabel.TITLE:
            parts.append(f"# {text}")
        elif label == DocItemLabel.SECTION_HEADER:
            parts.append(f"## {text}")
        else:
            parts.append(text)
    return "\n\n".join(parts)


# ── Docling local Granite VLM provider (Python API, no server) ─────────────────

def _login_huggingface_if_available() -> None:
    """Authenticate once if HF_TOKEN is present, then continue without auth on failure."""
    hf_token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGINGFACE_HUB_TOKEN")
    if not hf_token:
        return
    try:
        from huggingface_hub import login
        login(token=hf_token, add_to_git_credential=False)
    except Exception:
        pass


# ── Docling Granite VLM (full-page "Convert this page to docling." preset) ─────

def _pdf2md_docling_granite_local(pdf_bytes: bytes) -> str:
    """Run Docling's VLM pipeline with the Granite Docling preset."""
    import tempfile
    from pathlib import Path as _Path

    _login_huggingface_if_available()

    from docling.datamodel.base_models import InputFormat
    from docling.datamodel.pipeline_options import (
        AcceleratorDevice,
        AcceleratorOptions,
        VlmConvertOptions,
        VlmPipelineOptions,
    )
    from docling.document_converter import DocumentConverter, PdfFormatOption
    from docling.pipeline.vlm_pipeline import VlmPipeline

    vlm_options = VlmConvertOptions.from_preset("granite_docling")
    pipeline_options = VlmPipelineOptions(vlm_options=vlm_options)
    # AUTO picks CUDA/MPS if present, else CPU (slow — local is only for spot checks).
    pipeline_options.accelerator_options = AcceleratorOptions(device=AcceleratorDevice.AUTO)
    converter = DocumentConverter(
        format_options={
            InputFormat.PDF: PdfFormatOption(
                pipeline_cls=VlmPipeline,
                pipeline_options=pipeline_options,
            )
        }
    )

    with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as tmp:
        tmp.write(pdf_bytes)
        tmp_path = _Path(tmp.name)

    try:
        result = converter.convert(str(tmp_path))
        return _docling_doc_to_text(result.document)
    finally:
        tmp_path.unlink(missing_ok=True)


# ── Docling standard pipeline + formula/table enrichment (decodes formulas) ────

def _pdf2md_docling_standard_local(pdf_bytes: bytes) -> str:
    """Standard PDF pipeline with formula enrichment (formulas → LaTeX) and
    TableFormer (ACCURATE). Use this when the VLM leaves formulas undecoded —
    the CodeFormula model decodes them so `_latex_to_plain` can render them."""
    import tempfile
    from pathlib import Path as _Path

    _login_huggingface_if_available()

    from docling.datamodel.base_models import InputFormat
    from docling.datamodel.pipeline_options import (
        AcceleratorDevice,
        AcceleratorOptions,
        PdfPipelineOptions,
        TableFormerMode,
        TableStructureOptions,
    )
    from docling.document_converter import DocumentConverter, PdfFormatOption

    po = PdfPipelineOptions()
    po.do_ocr = False  # digital ACL PDFs already carry a text layer
    po.do_table_structure = True
    po.table_structure_options = TableStructureOptions(mode=TableFormerMode.ACCURATE)
    po.do_formula_enrichment = True  # decode formulas → LaTeX (CodeFormula model)
    po.generate_page_images = False
    po.accelerator_options = AcceleratorOptions(device=AcceleratorDevice.AUTO)
    converter = DocumentConverter(
        format_options={InputFormat.PDF: PdfFormatOption(pipeline_options=po)}
    )

    with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as tmp:
        tmp.write(pdf_bytes)
        tmp_path = _Path(tmp.name)
    try:
        result = converter.convert(str(tmp_path))
        return _docling_doc_to_text(result.document)
    finally:
        tmp_path.unlink(missing_ok=True)


async def _pdf2md_docling_server(pdf_bytes: bytes) -> str:
    """Send PDF to the Kaggle Docling server (Granite VLM) and get plain text back.

    The server linearizes the DoclingDocument to GT-style plain text (same
    template as test/parsed_pdfs), so the returned "markdown" is really the
    flattened text with only "# "/"## " headings preserved.
    """
    import httpx

    async def _post(url: str) -> "httpx.Response":
        async with httpx.AsyncClient(timeout=600) as client:
            return await client.post(
                f"{url}/parse",
                files={"file": ("paper.pdf", pdf_bytes, "application/pdf")},
            )

    base_url = pdf2md_backend._get_docling_url()
    resp = await _post(base_url)

    # Stale tunnel URL → re-discover via ntfy and retry once.
    if resp.status_code in (502, 503, 404):
        pdf2md_backend._docling_url_cache = None
        pdf2md_backend._cache_fetched_at = 0.0
        resp = await _post(pdf2md_backend._get_docling_url())

    resp.raise_for_status()
    return resp.json()["markdown"]


# ── Local PDF→Markdown provider (PyMuPDF) ─────────────────────────────────────

def _pdf2md_local(pdf_bytes: bytes) -> str:
    """
    Convert PDF to Markdown using PyMuPDF structured text extraction.
    Uses font size and bold flags to detect headings vs body text.
    """
    import fitz  # PyMuPDF
    import statistics

    doc = fitz.open(stream=pdf_bytes, filetype="pdf")

    # ── Pass 1: collect all font sizes to determine thresholds ──────────────
    all_sizes: list[float] = []
    for page in doc:
        for block in page.get_text("dict")["blocks"]:
            if block.get("type") != 0:
                continue
            for line in block.get("lines", []):
                for span in line.get("spans", []):
                    txt = span["text"].strip()
                    if txt and len(txt) > 2:
                        all_sizes.append(span["size"])

    if not all_sizes:
        doc.close()
        return ""

    body_size = statistics.median(all_sizes)
    # Thresholds calibrated for ACL two-column paper style:
    #   body ~10-11pt, section headings ~12pt (1.10x), title ~14pt (1.28x)
    h1_threshold = body_size * 1.25   # title / very large heading
    h2_threshold = body_size * 1.07   # section heading size baseline

    # Pattern for ACL-style section headings (content-based check).
    # Two tiers:
    #   1. Numbered headings: "1 Introduction", "2.1 Data" — very reliable.
    #   2. Keyword-prefix headings — keyword at start, short line (<= 60 chars).
    #      Matches "Introduction", "Introduction and Motivation", etc.
    #      Does NOT match long sentences that start with a keyword word.
    _numbered_heading_re = re.compile(
        r"^\d+(?:\.\d+)*\.?\s+[A-Z\d]", re.IGNORECASE
    )
    _keyword_heading_re = re.compile(
        r"^(Abstract|References?|Acknowledgments?|Related\s+Work"
        r"|Conclusion|Introduction|Appendix|Discussion"
        r"|Experiments?|Evaluation|Methodology|Datasets?)"
        r"(\b.{0,40})?$",   # allow up to 40 extra chars (subtitle words)
        re.IGNORECASE,
    )

    # ── Pass 2: build Markdown ───────────────────────────────────────────────
    lines_out: list[str] = []

    for page_num, page in enumerate(doc):
        page_height = page.rect.height
        blocks = page.get_text("dict")["blocks"]

        for block in blocks:
            if block.get("type") != 0:
                continue

            block_lines: list[str] = []
            for line in block.get("lines", []):
                line_text_parts: list[str] = []
                dominant_size = 0.0

                for span in line.get("spans", []):
                    txt = span["text"]
                    if not txt.strip():
                        line_text_parts.append(txt)
                        continue
                    sz = span["size"]
                    if sz > dominant_size:
                        dominant_size = sz
                    line_text_parts.append(txt)

                line_text = "".join(line_text_parts).strip()
                if not line_text:
                    continue

                # Skip very short noise (page numbers, line numbers < 4 chars)
                if len(line_text) <= 3 and line_text.isdigit():
                    continue
                # Skip header/footer areas (top/bottom 5% of page)
                block_y = block["bbox"][1]
                if block_y < page_height * 0.05 or block_y > page_height * 0.95:
                    continue

                # Determine heading level
                if dominant_size >= h1_threshold and page_num == 0:
                    # Title: large font on first page only, must be > 8 chars
                    if len(line_text) > 8:
                        line_text = f"# {line_text}"
                elif dominant_size >= h2_threshold and (
                    _numbered_heading_re.match(line_text)
                    or _keyword_heading_re.match(line_text)
                ):
                    # Section heading: larger font AND looks like an ACL section title.
                    # Combining both conditions avoids marking wrapped body lines as headings.
                    line_text = f"## {line_text}"

                block_lines.append(line_text)

            if block_lines:
                lines_out.append("\n".join(block_lines))

    doc.close()
    return "\n\n".join(lines_out)


# ── Markdown → PeerRead metadata conversion ────────────────────────────────────

def _split_sections(md: str) -> list[dict]:
    """Split Markdown into PeerRead-style [{heading, text}] blocks.

    Matches the ground-truth convention: every heading yields a section even
    when its body is empty (PeerRead keeps heading-only sections, e.g. a parent
    "3" before "3.1"). A leading headingless block is only kept if it has text.
    """
    lines = md.splitlines()
    sections: list[dict] = []
    current_heading: str | None = None
    current_lines: list[str] = []

    def _flush() -> None:
        text = "\n".join(current_lines).strip()
        # Headed sections are always emitted (text may be ""); a headingless
        # preamble is emitted only when it actually carries text.
        if current_heading is not None or text:
            sections.append({"heading": current_heading, "text": text})

    for line in lines:
        m = re.match(r"^(#{1,4})\s+(.+)", line)
        if m:
            _flush()
            current_heading = m.group(2).strip()
            current_lines = []
        else:
            current_lines.append(line)

    _flush()
    return sections


def _extract_title(md: str) -> str:
    # H1 first (PyMuPDF provider)
    for line in md.splitlines():
        m = re.match(r"^#\s+(.+)", line)
        if m:
            return m.group(1).strip()
    # Fallback: first H2 that doesn't look like a section heading (Docling uses H2 for titles)
    _section_re = re.compile(r"^(Abstract|Introduction|\d+[\s.])", re.IGNORECASE)
    for line in md.splitlines():
        m = re.match(r"^##\s+(.+)", line)
        if m:
            heading = m.group(1).strip()
            if not _section_re.match(heading):
                return heading
    return ""


def _extract_abstract(sections: list[dict]) -> str:
    for s in sections:
        if s["heading"] and re.search(r"\babstract\b", s["heading"], re.IGNORECASE):
            text = s["text"]
            # If the abstract section absorbed body text (> 2 000 chars) it means
            # subsequent section headings weren't detected. Return only the first
            # paragraph, which is the actual abstract.
            if len(text) > 2000:
                first_para = re.split(r"\n\n+", text)[0].strip()
                return first_para if len(first_para) > 50 else text[:2000]
            return text
    # Fallback: first headingless block long enough to be an abstract
    for s in sections:
        if not s["heading"] and len(s["text"]) > 150:
            return s["text"]
    return ""


def _extract_authors(md: str) -> list[str]:
    """
    Authors typically appear between the H1 title and the first section heading,
    as short lines of capitalized names (not affiliations or emails).
    """
    lines = md.splitlines()
    in_preamble = False
    preamble: list[str] = []

    for line in lines:
        if re.match(r"^#\s+", line):
            in_preamble = True
            continue
        if in_preamble:
            if re.match(r"^#{1,4}\s+", line):
                break
            preamble.append(line.strip())

    authors: list[str] = []
    for line in preamble:
        if not line:
            continue
        # Skip emails, URLs, affiliations
        if "@" in line or re.search(r"http|university|institute|department|lab\b|school", line, re.IGNORECASE):
            continue
        if len(line) > 100:
            continue
        # Must look like a name: starts capital, only letters/spaces/commas/hyphens/dots
        if re.match(r"^[A-Z][a-zA-Z\s,.\-]+$", line):
            parts = re.split(r",\s*|\s+and\s+", line)
            authors.extend(p.strip() for p in parts if p.strip())

    return authors


def _extract_year(md: str) -> int | None:
    # Only look in the preamble (before first section heading) to avoid
    # picking up years from citations in body text.
    preamble_end = len(md)
    for m in re.finditer(r"^#{1,4}\s+", md, re.MULTILINE):
        if m.start() > 0:
            preamble_end = m.start()
            break
    preamble = md[:min(preamble_end, 1500)]
    for m in re.finditer(r"\b(20[0-2]\d|19\d\d)\b", preamble):
        year = int(m.group(1))
        if 1990 <= year <= 2030:
            return year
    return None


# ── Reference parsing ──────────────────────────────────────────────────────────

def _last_name(author: str) -> str:
    """Last whitespace token of an author name, KEEPING any trailing period.

    GROBID's citeRegEx uses the raw token, so "Marianna Apidianaki." → "Apidianaki."
    and "Marco Baroni" → "Baroni" (the period only survives on the final author of
    the list, which is exactly how the ground truth builds it).
    """
    toks = author.split()
    return toks[-1] if toks else ""


def _build_cite_regex(authors: list[str], year: int | None) -> tuple[str, str]:
    """Return (citeRegEx, shortCiteRegEx) matching the GROBID author-year style."""
    if not authors or year is None:
        return "", ""

    def _name_pat(name: str) -> str:
        # Parsed last names may keep a trailing period from bibliography text.
        # Body citations usually do not, so make that period optional and escape
        # the rest of the name for regex safety.
        return re.escape(name.rstrip(".")).replace(r"\ ", r"\s+") + r"\.?"

    ln1 = _last_name(authors[0])
    if not ln1:
        return "", ""
    if len(authors) >= 3:
        short = rf"{_name_pat(ln1)}\s+et\s+al\.?"
    elif len(authors) == 2:
        # The second author is often mangled by PDF text extraction
        # (e.g. "M` arquez", "Pad´ o"). For mention detection, first-author
        # + "and <short name phrase>" + year is robust enough; the metric key
        # is first-author/year anyway.
        short = rf"{_name_pat(ln1)}\s+and\s+[A-Z][^,();]{{1,60}}?"
    else:
        short = _name_pat(ln1)

    # Match both parenthetical citations ("Name, 2016") and narrative citations
    # ("Name (2016)" or "Name (2015, 2016)"). The latter is common in PeerRead
    # bodies and was previously under-counted as a missing mention.
    year_pat = rf"{year}[a-z]?"
    parenthetical = rf"{short},?\s+{year_pat}"
    narrative = rf"{short}\s*\([^)]*\b{year_pat}\b[^)]*\)"
    return rf"(?:{parenthetical}|{narrative})", short


def _parse_reference(text: str) -> dict | None:
    """Parse one author-year reference entry → PeerRead reference dict.

    Expected ACL layout: "Authors. Year. Title. Venue." (also tolerates a
    leading "[N]"/"N."/bullet and a quoted title).
    """
    text = re.sub(r"^[\[(]?\d+[\]).]?\s*", "", text).strip()  # leading [N] / N.
    text = re.sub(r"^[-*•·]\s*", "", text).strip()            # leading bullet
    if len(text) < 10:
        return None

    # Author/year boundary = first 19xx/20xx, preferring one followed by a period.
    year_m = (
        re.search(r"\b((?:19|20)\d\d)[a-z]?\.", text)
        or re.search(r"\b((?:19|20)\d\d)[a-z]?\b", text)
    )
    if not year_m:
        return None
    year = int(year_m.group(1))

    authors_str = text[: year_m.start()].rstrip()
    rest = text[year_m.end():].strip()

    authors: list[str] = []
    if authors_str:
        parts = re.split(r",\s*(?:and\s+)?|\s+and\s+", authors_str)
        authors = [p.strip() for p in parts if p.strip() and len(p.strip()) < 60]

    # Title / venue from the remainder ("Title. Venue").
    title = venue = ""
    qm = re.search(r'["“](.+?)["”]', rest)
    if qm:
        title = qm.group(1).strip()
        venue = rest[qm.end():].strip(' .,')
    elif rest:
        tparts = rest.split(". ", 1)
        title = tparts[0].strip().rstrip(".")
        venue = tparts[1].strip() if len(tparts) > 1 else ""

    cite_regex, short_cite = _build_cite_regex(authors, year)
    return {
        "title": title,
        "author": authors,
        "venue": venue,
        "citeRegEx": cite_regex,
        "shortCiteRegEx": short_cite,
        "year": year,
    }


def _split_ref_entries(block: str) -> list[str]:
    """Split a References block into per-entry strings.

    Docling usually emits each bibliography item as its own block (blank-line
    separated). If that doesn't hold, fall back to grouping wrapped lines: a new
    entry begins at an author-like line once the buffer already carries a year.
    """
    block = block.strip()
    year_re = re.compile(r"\b(?:19|20)\d\d")

    paras = [p.strip() for p in re.split(r"\n\s*\n", block) if p.strip()]
    if len(paras) >= 2 and sum(1 for p in paras if year_re.search(p)) >= len(paras) * 0.6:
        return paras

    lines = [ln.strip() for ln in block.splitlines() if ln.strip()]
    start_re = re.compile(r"^[A-Z][A-Za-z.'\-]+(?:,|\s+[A-Z])")  # "Surname," / "Surname Initial"
    entries: list[str] = []
    buf: list[str] = []
    for ln in lines:
        if buf and year_re.search(" ".join(buf)) and start_re.match(ln):
            entries.append(" ".join(buf))
            buf = [ln]
        else:
            buf.append(ln)
    if buf:
        entries.append(" ".join(buf))
    return entries


def _parse_references(md: str) -> list[dict]:
    """Extract and parse the References / Bibliography section."""
    m = re.search(
        r"^#{1,4}\s+(?:References?|Bibliography)\s*$", md, re.MULTILINE | re.IGNORECASE
    )
    if not m:
        return []

    ref_block = md[m.end():]
    next_h = re.search(r"^#{1,4}\s+", ref_block, re.MULTILINE)  # cut at next heading
    if next_h:
        ref_block = ref_block[: next_h.start()]

    refs = []
    for entry in _split_ref_entries(ref_block):
        r = _parse_reference(entry)
        if r:
            refs.append(r)
    return refs


def _sentence_context(text: str, start: int, end: int) -> tuple[str, int, int]:
    """Sentence containing [start, end), plus the match offsets within it.

    GROBID's referenceMentions use sentence-bounded context (not a fixed window).
    """
    prev = list(re.finditer(r"[.!?]\s+", text[:start]))
    ctx_start = prev[-1].end() if prev else 0
    nxt = re.search(r"[.!?]\s+", text[end:])
    ctx_end = end + nxt.start() + 1 if nxt else len(text)

    ctx = text[ctx_start:ctx_end]
    lead = len(ctx) - len(ctx.lstrip())
    ctx = ctx.strip()
    return ctx, start - ctx_start - lead, end - ctx_start - lead


def _extract_reference_mentions(sections: list[dict], references: list[dict]) -> list[dict]:
    """Scan body text for author-year citations and build referenceMentions.

    Each reference's citeRegEx is matched against every body section (the
    References section itself is skipped). Mentions are emitted in document order
    with sentence-bounded context + offsets, matching the PeerRead schema.
    """
    compiled: list[tuple[int, tuple[str, int] | None, "re.Pattern[str]"]] = []
    for i, r in enumerate(references):
        pat = r.get("citeRegEx") or ""
        if not pat:
            continue
        authors = r.get("author") or []
        toks = authors[0].split() if authors else []
        ln = toks[-1].strip(".,").lower() if toks else ""
        year = r.get("year")
        key = (ln, year) if ln and isinstance(year, int) else None
        try:
            compiled.append((i, key, re.compile(pat)))
        except re.error:
            continue

    mentions: list[dict] = []
    for section in sections:
        heading = (section.get("heading") or "").strip()
        if re.match(r"references?$", heading, re.IGNORECASE):
            continue
        text = section.get("text", "")
        if not text:
            continue

        found: list[tuple[int, int, int, tuple[str, int] | None]] = []
        seen_spans: set[tuple[int, int, tuple[str, int] | None]] = set()
        for ref_id, key, rx in compiled:
            for mm in rx.finditer(text):
                span_key = (mm.start(), mm.end(), key)
                if span_key in seen_spans:
                    continue
                seen_spans.add(span_key)
                found.append((mm.start(), mm.end(), ref_id, key))
        found.sort()  # document order

        for s_, e_, ref_id, _key in found:
            ctx, off_s, off_e = _sentence_context(text, s_, e_)
            mentions.append({
                "referenceID": ref_id,
                "context": ctx,
                "startOffset": off_s,
                "endOffset": off_e,
            })

    return mentions


def _refresh_reference_regexes(references: list[dict]) -> list[dict]:
    """Rebuild citation regexes with the current parser logic."""
    refreshed: list[dict] = []
    for ref in references:
        r = dict(ref)
        cite_regex, short_cite = _build_cite_regex(r.get("author") or [], r.get("year"))
        if cite_regex:
            r["citeRegEx"] = cite_regex
            r["shortCiteRegEx"] = short_cite
        refreshed.append(r)
    return refreshed


def _sections_for_mention_eval(meta: dict) -> list[dict]:
    """Sections used when rebuilding mentions from saved PeerRead-style JSON.

    `md_to_peerread` extracts mentions before dropping Abstract from `sections`,
    but saved JSON keeps the abstract separately as `abstractText`. Rebuild from
    both sides with the same source text so the metric is symmetric.
    """
    sections: list[dict] = []
    abstract = (meta.get("abstractText") or "").strip()
    if abstract:
        sections.append({"heading": "Abstract", "text": abstract})
    sections.extend(meta.get("sections", []))
    return sections


def _normalize_mentions_for_eval(meta: dict) -> dict:
    """Rebuild references' regexes and mentions with the current parser logic."""
    out = dict(meta)
    out["references"] = _refresh_reference_regexes(out.get("references", []))
    out["referenceMentions"] = _extract_reference_mentions(
        _sections_for_mention_eval(out), out["references"]
    )
    return out


# ── Top-level converter ────────────────────────────────────────────────────────

def md_to_peerread(md: str, paper_id: str, provider: str = "docling") -> dict:
    """Convert Markdown (from any provider) to PeerRead metadata JSON."""
    sections = _split_sections(md)
    title = _extract_title(md)
    abstract = _extract_abstract(sections)
    authors = _extract_authors(md)
    year = _extract_year(md)
    references = _parse_references(md)
    mentions = _extract_reference_mentions(sections, references)

    # Match the ground-truth `sections` convention: it excludes the title block,
    # the Abstract (lives only in abstractText) and the References bibliography
    # (lives only in references[]). Appendices after the bibliography are kept.
    def _is_excluded(heading: str | None) -> bool:
        if heading is None:
            return False
        h = heading.strip()
        if h == title:
            return True
        if re.search(r"\babstract\b", h, re.IGNORECASE):
            return True
        if re.match(r"references?$", h, re.IGNORECASE):
            return True
        return False

    body_sections = [s for s in sections if not _is_excluded(s.get("heading"))]

    return {
        "name": f"{paper_id}.pdf",
        "metadata": {
            "source": provider,
            "title": title,
            "authors": authors,
            "emails": [],
            "sections": body_sections,
            "references": references,
            "referenceMentions": mentions,
            "year": year,
            "abstractText": abstract,
            "creator": provider.capitalize(),
        },
    }


# ── OCR runner ─────────────────────────────────────────────────────────────────

async def _get_markdown(pdf_bytes: bytes, provider: str) -> str:
    if provider == "local":
        return _pdf2md_local(pdf_bytes)
    if provider == "docling-granite":
        return await _pdf2md_docling_server(pdf_bytes)
    if provider == "docling-granite-local":
        return _pdf2md_docling_granite_local(pdf_bytes)
    if provider == "docling-standard-local":
        return _pdf2md_docling_standard_local(pdf_bytes)
    raise ValueError(f"Unknown provider: {provider!r}")


async def process_pdf(
    pdf_path: Path, out_dir: Path, provider: str, overwrite: bool = False
) -> None:
    paper_id = pdf_path.stem
    out_path = out_dir / f"{paper_id}.pdf.json"

    if out_path.exists() and not overwrite:
        print(f"[skip] {paper_id}  (already done — use --overwrite to redo)")
        return

    label = {
        "local": "locally (PyMuPDF)",
        "docling-granite": "Docling Granite VLM pipeline (Kaggle server)",
        "docling-granite-local": "Docling Granite VLM pipeline (local)",
        "docling-standard-local": "Docling standard pipeline + formula enrichment (local)",
    }.get(provider, provider)
    print(f"[ .. ] {paper_id}  processing via {label} ...", end="", flush=True)
    try:
        md = await _get_markdown(pdf_path.read_bytes(), provider)
        result = md_to_peerread(md, paper_id, provider)
        out_path.write_text(json.dumps(result, ensure_ascii=False, indent=2), encoding="utf-8")
        n_sec = len(result["metadata"]["sections"])
        n_ref = len(result["metadata"]["references"])
        print(f"\r[ ok ] {paper_id}  {n_sec} sections  {n_ref} refs  -> {out_path.name}")
    except Exception as exc:
        import traceback
        print(f"\r[ERR] {paper_id}  {exc}")
        traceback.print_exc()


async def run_ocr(pdf_dir: Path, out_dir: Path, provider: str, overwrite: bool) -> None:
    # Resolve to absolute path before any library call that might change CWD
    out_dir = out_dir.resolve()
    pdf_dir = pdf_dir.resolve()
    out_dir.mkdir(parents=True, exist_ok=True)
    pdfs = sorted(pdf_dir.glob("*.pdf"))
    if not pdfs:
        sys.exit(f"No PDFs found in {pdf_dir}")
    print(f"Processing {len(pdfs)} PDFs from {pdf_dir}  [provider={provider}]\n")
    for pdf in pdfs:
        await process_pdf(pdf, out_dir, provider, overwrite)
    print(f"\nDone. Output in {out_dir}/")


# ── Comparison report ──────────────────────────────────────────────────────────

def _normalize_for_eval(text: str) -> str:
    """Canonicalize text before similarity scoring, applied to BOTH GT and pred.

    The ground truth (GROBID) carries noise that Docling cleanly omits — margin
    line numbers, page numbers, mojibake (σ→Ï, −→â1) and odd formula spacing.
    Comparing raw text unfairly penalizes the cleaner output, so we fold both
    sides to a common form: drop standalone numeric lines, fold non-ASCII
    (mojibake and real unicode symbols alike), lowercase, collapse whitespace.
    """
    lines = [
        ln for ln in text.splitlines()
        if not re.fullmatch(r"\s*\d{1,4}[.)]?\s*", ln)  # pure margin/page numbers
    ]
    s = " ".join(lines)
    s = unicodedata.normalize("NFKD", s).encode("ascii", "ignore").decode("ascii")
    s = re.sub(r"\s+", " ", s).strip().lower()
    return s


def _body_text(meta: dict) -> str:
    return "\n".join(s.get("text", "") for s in meta.get("sections", []))


def _ref_keys(meta: dict) -> set[tuple[str, int]]:
    """(first-author last name, year) keys — robust to formatting differences."""
    keys: set[tuple[str, int]] = set()
    for r in meta.get("references", []):
        authors = r.get("author") or []
        toks = authors[0].split() if authors else []
        ln = toks[-1].strip(".,").lower() if toks else ""
        year = r.get("year")
        if ln and isinstance(year, int):
            keys.add((ln, year))
    return keys


def _loose(text: str) -> str:
    """Space-insensitive canonical form (applied to BOTH sides): drop margin/page
    numbers, fold non-ASCII (GT mojibake σ→Ï), lowercase, and REMOVE ALL spaces.

    OCR often inserts spaces between characters (e.g. "h i -1", "s ( c ) k"); the
    ground truth spaces them differently. Removing whitespace makes scoring lenient
    about spacing so it reflects actual character content, not layout.
    """
    lines = [
        ln for ln in text.splitlines()
        if not re.fullmatch(r"\s*\d{1,4}[.)]?\s*", ln)
    ]
    s = " ".join(lines)
    s = unicodedata.normalize("NFKD", s).encode("ascii", "ignore").decode("ascii").lower()
    return re.sub(r"\s+", "", s)


def _sim_loose(a: str, b: str) -> float:
    """Character similarity (↑, 0–1) ignoring whitespace differences."""
    la, lb = _loose(a), _loose(b)
    if not la and not lb:
        return 1.0
    return SequenceMatcher(None, la, lb, autojunk=False).ratio()


def _tokens_for_eval(text: str) -> list[str]:
    """Word/number tokens after the same loose normalization used for scoring."""
    return re.findall(r"[a-z0-9]+(?:[.-][a-z0-9]+)*", _normalize_for_eval(text))


def _multiset_prf(gt_items: list, pred_items: list) -> tuple[float, float, float]:
    """Precision/recall/F1 over multisets (both empty ⇒ perfect)."""
    from collections import Counter
    g, p = Counter(gt_items), Counter(pred_items)
    if not g and not p:
        return 1.0, 1.0, 1.0
    inter = sum((g & p).values())
    rec = inter / sum(g.values()) if g else 0.0
    prec = inter / sum(p.values()) if p else 0.0
    f1 = 2 * prec * rec / (prec + rec) if (prec + rec) else 0.0
    return prec, rec, f1


def _classify_lines(text: str) -> tuple[list[str], list[str]]:
    """Split body text into (prose lines, table-ish lines).

    A table-ish line is number-dense (≥3 numeric tokens and ≥40% of tokens
    numeric) — GROBID flattens tables into exactly this shape.
    """
    prose, table = [], []
    for ln in text.splitlines():
        s = ln.strip()
        if not s or re.fullmatch(r"\d{1,4}[.)]?", s):
            continue
        toks = s.split()
        nums = [t for t in toks if re.search(r"\d", t)]
        if toks and len(nums) >= 3 and len(nums) >= 0.4 * len(toks):
            table.append(s)
        else:
            prose.append(s)
    return prose, table


_CLEAN_CHAR_RE = re.compile(r"[A-Za-z0-9 .,%()\[\];:/+\-#|~]")


def _is_junk_line(line: str) -> bool:
    """True for a GROBID line that is really a mis-OCR'd figure/CJK block.

    PeerRead's ground truth sometimes flattens a figure, CJK caption or score
    glyph into the body as mojibake dominated by punctuation ($ & ! " 架 …).
    These leak into the table/formula buckets and unfairly drag scores down, so
    numeric/formula scoring skips them. Threshold: a genuine data row is ≥80%
    "clean" chars (letters, digits, %, brackets, common punctuation); mojibake
    rows fall below it.
    """
    s = line.strip()
    if not s:
        return True
    clean = sum(1 for c in s if _CLEAN_CHAR_RE.match(c))
    return clean / len(s) < 0.80


def _data_numbers(text: str) -> list[str]:
    """Meaningful numeric values: decimals, percentages, or ≥2-digit/comma counts.

    Drops lone single digits — GROBID renders checkmarks/daggers as "7" and
    section refs as "3", which are noise rather than reported table data.
    """
    out: list[str] = []
    for t in re.findall(r"\d[\d.,%]*", text):
        core = t.rstrip(".,%")
        if "." in core or "%" in t or "," in core or len(core) >= 2:
            out.append(core)
    return out


def _heading_seq(meta: dict) -> list[str]:
    headings = []
    for s in meta.get("sections", []):
        if not s.get("heading"):
            continue
        h = _normalize_for_eval(s["heading"])
        # Output can keep submission boilerplate that GT dropped. That should not
        # count against reading-order preservation.
        if h and h != "anonymous acl submission":
            headings.append(h)
    return headings


def _lcs_len(a: list, b: list) -> int:
    if not a or not b:
        return 0
    prev = [0] * (len(b) + 1)
    for x in a:
        cur = [0] * (len(b) + 1)
        for j, y in enumerate(b):
            cur[j + 1] = prev[j] + 1 if x == y else max(prev[j + 1], cur[j])
        prev = cur
    return prev[-1]


def _mention_keys(meta: dict) -> list[tuple[str, int]]:
    """Resolve each referenceMention to its (last-name, year) via referenceID."""
    refs = meta.get("references", [])
    keys: list[tuple[str, int]] = []
    for mn in meta.get("referenceMentions", []):
        rid = mn.get("referenceID")
        if isinstance(rid, int) and 0 <= rid < len(refs):
            authors = refs[rid].get("author") or []
            toks = authors[0].split() if authors else []
            ln = toks[-1].strip(".,").lower() if toks else ""
            year = refs[rid].get("year")
            if ln and isinstance(year, int):
                keys.append((ln, year))
    return keys


def _component_scores(gt: dict, pred: dict) -> dict:
    """Per-component metrics for one paper (see legend in run_compare)."""
    gt_body, pred_body = _body_text(gt), _body_text(pred)
    gt_prose, gt_table = _classify_lines(gt_body)
    pr_prose, pr_table = _classify_lines(pred_body)

    gt_prose_text = "\n".join(gt_prose)
    pr_prose_text = "\n".join(pr_prose)
    text_sim = _sim_loose(gt_prose_text, pr_prose_text)                 # Text char diagnostic ↑
    text_p, text_r, text_f = _multiset_prf(
        _tokens_for_eval(gt_prose_text), _tokens_for_eval(pr_prose_text)
    )
    # Table: coverage of the GROBID table's DATA VALUES in the prediction.
    # GROBID and Docling flatten tables differently (space- vs dot-joined), so
    # classifying pred lines independently is unreliable — it scored a paper
    # 0.000 although every number was captured. Instead take the data values
    # GROBID placed in table-ish rows (minus mojibake figure rows) and measure
    # how many appear ANYWHERE in the prediction body: distinct-value recall,
    # robust to flatten/spacing and to GROBID's spurious duplicated rows.
    gt_table_clean = [l for l in gt_table if not _is_junk_line(l)]
    pr_table_clean = [l for l in pr_table if not _is_junk_line(l)]
    table_sim = _sim_loose("\n".join(gt_table_clean), "\n".join(pr_table_clean))  # diagnostic ↑
    gt_dnums = set(_data_numbers("\n".join(gt_table_clean)))
    pred_dnums_all = set(_data_numbers(pred_body))
    tbl_num_cov = len(gt_dnums & pred_dnums_all) / len(gt_dnums) if gt_dnums else 1.0
    # multiset P/R/F on the table buckets, kept as diagnostics only
    tbl_num_p, tbl_num_r, tbl_num_f = _multiset_prf(
        _data_numbers("\n".join(gt_table_clean)), _data_numbers("\n".join(pr_table_clean))
    )
    # Formula: equation recall via "=" count, on non-junk lines only so a mojibake
    # "=" inside a mis-OCR'd figure row can't inflate the GT denominator.
    eg = sum(ln.count("=") for ln in gt_body.splitlines() if not _is_junk_line(ln))
    ep = sum(ln.count("=") for ln in pred_body.splitlines() if not _is_junk_line(ln))
    formula_rec = min(ep, eg) / eg if eg else 1.0
    gh, ph = _heading_seq(gt), _heading_seq(pred)                        # Order LCS ↑
    order = _lcs_len(gh, ph) / len(gh) if gh else (1.0 if not ph else 0.0)
    ref_p, ref_r, ref_f = _multiset_prf(list(_ref_keys(gt)), list(_ref_keys(pred)))
    men_p, men_r, men_f = _multiset_prf(_mention_keys(gt), _mention_keys(pred))

    return {
        "text_sim": text_sim, "text_p": text_p, "text_r": text_r, "text_f": text_f,
        "table_sim": table_sim,
        "tbl_num_cov": tbl_num_cov,
        "tbl_num_p": tbl_num_p, "tbl_num_r": tbl_num_r, "tbl_num_f": tbl_num_f,
        "formula": formula_rec, "order": order,
        "ref_p": ref_p, "ref_r": ref_r, "ref_f": ref_f,
        "men_p": men_p, "men_r": men_r, "men_f": men_f,
    }


def run_compare(gt_dir: Path, pred_dir: Path) -> None:
    gt_files = {p.name: p for p in gt_dir.glob("*.pdf.json")}
    pred_files = {p.name: p for p in pred_dir.glob("*.pdf.json")}
    common = sorted(gt_files.keys() & pred_files.keys())

    if not common:
        sys.exit(f"No matching files between {gt_dir} and {pred_dir}")

    rows: list[tuple[str, dict]] = []
    for name in common:
        gt = json.loads(gt_files[name].read_text(encoding="utf-8"))["metadata"]
        pred = json.loads(pred_files[name].read_text(encoding="utf-8"))["metadata"]
        # Apply parser improvements symmetrically without forcing a full Docling
        # re-run. This avoids comparing normalized pred mentions against older
        # GROBID mention extraction quirks from the saved GT JSON.
        gt = _normalize_mentions_for_eval(gt)
        pred = _normalize_mentions_for_eval(pred)
        rows.append((name.split(".")[0], _component_scores(gt, pred)))

    n = len(rows)
    avg = lambda k: sum(s[k] for _, s in rows) / n

    print(f"Comparing {n} paper(s)   (all scores 0–1, HIGHER = better, space-insensitive)\n")

    # ── Headline: one easy score per data type ───────────────────────────────
    sh = (f"{'ID':<6} {'Text':>6} {'Table':>6} {'Formula':>8} "
          f"{'Order':>6} {'Refs':>6} {'Mentions':>9}")
    print(sh)
    print("-" * len(sh))
    for pid, s in rows:
        print(f"{pid:<6} {s['text_r']:>6.3f} {s['tbl_num_cov']:>6.3f} {s['formula']:>8.3f} "
              f"{s['order']:>6.3f} {s['ref_r']:>6.3f} {s['men_r']:>9.3f}")
    print("-" * len(sh))
    print(f"{'AVG':<6} {avg('text_r'):>6.3f} {avg('tbl_num_cov'):>6.3f} {avg('formula'):>8.3f} "
          f"{avg('order'):>6.3f} {avg('ref_r'):>6.3f} {avg('men_r'):>9.3f}")

    # ── Precision / Recall detail ────────────────────────────────────────────
    dh = (f"{'ID':<6} {'TxtP':>6} {'TxtF':>6} {'TxtChr':>6} "
          f"{'TblP':>6} {'TblR':>6} {'TblF':>6} {'TblTxt':>6} "
          f"{'RefP':>6} {'RefR':>6} {'MenP':>6} {'MenR':>6}")
    print("\n=== Detail (text/table diagnostics, reference & mention P/R) ===")
    print(dh)
    print("-" * len(dh))
    for pid, s in rows:
        print(f"{pid:<6} {s['text_p']:>6.3f} {s['text_f']:>6.3f} {s['text_sim']:>6.3f} "
              f"{s['tbl_num_p']:>6.3f} {s['tbl_num_r']:>6.3f} "
              f"{s['tbl_num_f']:>6.3f} {s['table_sim']:>6.3f} "
              f"{s['ref_p']:>6.3f} {s['ref_r']:>6.3f} {s['men_p']:>6.3f} {s['men_r']:>6.3f}")
    print("-" * len(dh))
    print(f"{'AVG':<6} {avg('text_p'):>6.3f} {avg('text_f'):>6.3f} {avg('text_sim'):>6.3f} "
          f"{avg('tbl_num_p'):>6.3f} {avg('tbl_num_r'):>6.3f} "
          f"{avg('tbl_num_f'):>6.3f} {avg('table_sim'):>6.3f} "
          f"{avg('ref_p'):>6.3f} {avg('ref_r'):>6.3f} {avg('men_p'):>6.3f} {avg('men_r'):>6.3f}")

    print("\nText     ↑ : GT-token recall on prose (does not punish useful extra extraction).")
    print("Table    ↑ : distinct GT table data-values found in pred (junk rows filtered).")
    print("Formula  ↑ : formula recall (decoded equations captured vs GT).")
    print("Order    ↑ : GT heading-order recall (LCS / #GT headings).")
    print("Refs     ↑ : reference recall by (last-name, year).")
    print("Mentions ↑ : citation→reference linking recall.")
    print("TxtChr/TblTxt ↑ : old char-similarity diagnostics only.")


# ── Entry point ────────────────────────────────────────────────────────────────

def main() -> None:
    parser = argparse.ArgumentParser(
        description="Test OCR against PeerRead ACL-2017 test set"
    )
    parser.add_argument("--pdf-dir", default="test/pdfs", help="Input PDF directory")
    parser.add_argument("--out-dir", default="test/docling_output", help="Output JSON directory")
    parser.add_argument(
        "--provider",
        default="docling-granite",
        choices=[
            "docling-granite",        # Granite VLM via Kaggle server
            "docling-granite-local",  # Granite VLM locally (needs GPU; slow on CPU)
            "docling-standard-local", # standard pipeline + formula enrichment (decodes formulas)
            "local",                  # PyMuPDF baseline
        ],
        help="docling-standard-local=decodes formulas (recommended), docling-granite=VLM via Kaggle, local=PyMuPDF baseline",
    )
    parser.add_argument("--overwrite", action="store_true", help="Re-process already-done files")
    parser.add_argument(
        "--compare",
        action="store_true",
        help="Compare test/docling_output/ vs test/parsed_pdfs/ (ground truth)",
    )
    parser.add_argument(
        "--gt-dir",
        default="test/parsed_pdfs",
        help="Ground truth directory (for --compare)",
    )
    args = parser.parse_args()

    if args.compare:
        run_compare(Path(args.gt_dir), Path(args.out_dir))
    else:
        asyncio.run(run_ocr(Path(args.pdf_dir), Path(args.out_dir), args.provider, args.overwrite))


if __name__ == "__main__":
    main()