codex-py/codex/parsing/tex.py

"""LaTeX parsing utilities.

Provides helpers to extract sections, chunk text, and convert LaTeX to plain
readable prose by stripping markup that is not useful for NLP/embedding.
"""

from __future__ import annotations

import re

# ---------------------------------------------------------------------------
# Internal helpers
# ---------------------------------------------------------------------------

_SECTION_RE = re.compile(
    r"\\(?:sub)*section\*?\s*\{([^}]*)\}",
    re.DOTALL,
)

# \input{file} / \include{file} — arXiv papers routinely split the body across
# files, so these must be inlined before parsing (R-C multi-file flattening).
_INPUT_RE = re.compile(r"\\(?:input|include)\s*\{([^}]+)\}")

# Patterns for LaTeX noise removal (applied in order).
_COMMENT_RE = re.compile(r"%[^\n]*")
_CITE_RE = re.compile(r"\\cite\*?\{[^}]*\}")
_LABEL_RE = re.compile(r"\\label\{[^}]*\}")
_REF_RE = re.compile(r"\\ref\{[^}]*\}")
# Display math: $$...$$  (before single $ to avoid greedy mismatch)
_DISPLAY_DOLLAR_RE = re.compile(r"\$\$.*?\$\$", re.DOTALL)
# Inline math: $...$
_INLINE_MATH_RE = re.compile(r"\$[^$\n]*?\$", re.DOTALL)
# \begin{equation}...\end{equation}
_ENV_EQUATION_RE = re.compile(
    r"\\begin\{equation\*?\}.*?\\end\{equation\*?\}",
    re.DOTALL,
)
# \begin{align}...\end{align}  (covers align, align*, aligned, …)
_ENV_ALIGN_RE = re.compile(
    r"\\begin\{align[^}]*\}.*?\\end\{align[^}]*\}",
    re.DOTALL,
)
# Collapse excess whitespace
_WHITESPACE_RE = re.compile(r"[ \t]+")
_BLANK_LINES_RE = re.compile(r"\n{3,}")


def _clean_latex(text: str) -> str:
    """Strip LaTeX markup and return readable prose."""
    text = _COMMENT_RE.sub("", text)
    text = _ENV_EQUATION_RE.sub("", text)
    text = _ENV_ALIGN_RE.sub("", text)
    text = _DISPLAY_DOLLAR_RE.sub("", text)
    text = _INLINE_MATH_RE.sub("", text)
    text = _CITE_RE.sub("", text)
    text = _LABEL_RE.sub("", text)
    text = _REF_RE.sub("", text)
    text = _WHITESPACE_RE.sub(" ", text)
    text = _BLANK_LINES_RE.sub("\n\n", text)
    return text.strip()


# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------


def extract_sections(latex: str) -> list[tuple[str, str]]:
    """Split *latex* on \\section / \\subsection boundaries.

    Returns a list of ``(title, cleaned_text)`` tuples, one per section.
    Text between the document start and the first section command is
    discarded (preamble / abstract handling is out of scope).
    """
    # Find all section positions
    matches = list(_SECTION_RE.finditer(latex))
    if not matches:
        return []

    sections: list[tuple[str, str]] = []
    for idx, match in enumerate(matches):
        title = match.group(1).strip()
        body_start = match.end()
        body_end = matches[idx + 1].start() if idx + 1 < len(matches) else len(latex)
        body = latex[body_start:body_end]
        cleaned = _clean_latex(body)
        sections.append((title, cleaned))

    return sections


def chunk_text(text: str, size: int = 512, overlap: int = 64) -> list[str]:
    """Split *text* into overlapping word-based chunks.

    Parameters
    ----------
    text:
        Plain text to chunk (not raw LaTeX).
    size:
        Target number of words per chunk.
    overlap:
        Number of words to carry over into the next chunk.

    Each chunk is at most ``size + overlap`` words.  Boundaries are snapped
    to the nearest ``". "`` (sentence end) within ±20 words of the nominal
    boundary when possible.
    """
    words = text.split()
    if not words:
        return []

    snap_window = 20
    chunks: list[str] = []
    start = 0

    while start < len(words):
        end = min(start + size, len(words))

        # Try to snap *end* to a sentence boundary within ±snap_window words.
        if end < len(words):
            best = end
            # Build a small search window
            lo = max(start + 1, end - snap_window)
            hi = min(len(words), end + snap_window + 1)
            # Prefer the closest sentence-end ". " to *end*
            for offset in range(0, snap_window + 1):
                for candidate in (end - offset, end + offset):
                    if (
                        lo <= candidate < hi
                        and candidate > start
                        and words[candidate - 1].endswith(".")
                    ):
                        # Sentence boundary: word at (candidate-1) ends with ".".
                        best = candidate
                        break
                if best != end:
                    break
            end = best

        chunk_words = words[start:end]
        chunks.append(" ".join(chunk_words))

        if end >= len(words):
            break

        # Next chunk starts *overlap* words before *end*.
        start = max(start + 1, end - overlap)

    return chunks


def latex_to_text(latex: str) -> str:
    """Convert a LaTeX document to plain text.

    Extracts all sections, cleans each one, and joins them with a blank line.
    If no section commands are found, the whole document is cleaned and
    returned as a single block.
    """
    sections = extract_sections(latex)
    if sections:
        return "\n\n".join(body for _, body in sections)
    # Fallback: no section structure — clean the whole string.
    return _clean_latex(latex)


def _norm_texkey(name: str) -> str:
    """Normalize a tex file name / ``\\input`` argument to a comparable key.

    Strips a leading ``./``, any directory-irrelevant whitespace, and a trailing
    ``.tex`` so ``\\input{sections/intro}`` matches the archive member
    ``sections/intro.tex``.
    """
    n = name.strip().lstrip("./")
    return n[:-4] if n.endswith(".tex") else n


def flatten_inputs(primary: str, files: dict[str, str]) -> str:
    """Recursively inline ``\\input``/``\\include`` directives (R-C).

    *files* maps each source file's archive name to its contents (keys may carry
    a ``.tex`` suffix or a ``./`` prefix — they are normalized internally). arXiv
    multi-file projects leave the ``\\documentclass`` file as little more than a
    skeleton of ``\\input`` lines, so the body must be assembled before parsing.
    Unresolved references are dropped; recursion is capped at depth 10 to guard
    against include cycles.
    """
    norm = {_norm_texkey(k): v for k, v in files.items()}

    def _expand(text: str, depth: int) -> str:
        if depth > 10:
            return text

        def _repl(match: re.Match[str]) -> str:
            content = norm.get(_norm_texkey(match.group(1)))
            return _expand(content, depth + 1) if content is not None else ""

        return _INPUT_RE.sub(_repl, text)

    return _expand(primary, 0)


def chunk_sections(latex: str, size: int = 512, overlap: int = 64) -> list[tuple[str | None, str]]:
    """Chunk LaTeX *within* each section, pairing every chunk with its title.

    Unlike :func:`latex_to_text` (which discards titles and flattens the whole
    document into one block before word-window chunking), this keeps chunks from
    spanning section boundaries and remembers which ``\\section`` each came from.
    The caller can then label the stored ``section`` column from the real heading
    instead of mostly ``body`` (R-C / audit R-12). Falls back to title-less
    chunks of the whole cleaned document when there are no section commands.
    """
    sections = extract_sections(latex)
    if not sections:
        return [(None, chunk) for chunk in chunk_text(_clean_latex(latex), size, overlap)]
    pairs: list[tuple[str | None, str]] = []
    for title, body in sections:
        for chunk in chunk_text(body, size, overlap):
            pairs.append((title or None, chunk))
    return pairs