lllin000_PaperForge/paperforge/worker/ocr_objects.py

from __future__ import annotations

import contextlib
import re
from collections.abc import Callable
from pathlib import Path
from typing import Any

from paperforge.worker.ocr_math import normalize_ocr_math_text


def render_figure_object_markdown(figure: dict[str, Any]) -> str:
    caption = figure.get("caption", "")
    image_relpath = figure.get("image_relpath", "")

    # Extract figure number for the title
    figure_id = figure.get("figure_id", "")
    if figure_id.startswith("unresolved_cluster_") or figure_id.startswith("cluster_"):
        label = "Unresolved Figure Candidate"
    elif figure_id and not figure_id.startswith("orphan_"):
        m = re.search(r"\d+", figure_id)
        num = str(int(m.group())) if m else figure_id
        label = f"Figure {num}"
    else:
        label = "Orphan Media"

    was_cropped = figure.get("was_cropped", True)
    parts = [f"# {label}", ""]
    if was_cropped and image_relpath:
        parts.append(f"![](../../{image_relpath})")
    parts.append("")
    if caption:
        parts.append("## Legend")
        parts.append(normalize_ocr_math_text(caption))
    if figure.get("page"):
        parts.append("")
        parts.append(f"*Page {figure['page']}*")
    parts.append("")
    parts.append("---")

    return "\n".join(parts)


def render_table_object_markdown(table: dict[str, Any]) -> str:
    caption = table.get("caption", "")
    image_relpath = table.get("image_relpath", "")

    formal_num = table.get("formal_table_number")
    if formal_num is not None:
        label = f"Table {formal_num}"
    else:
        table_id_raw = table.get("table_id", "unknown")
        m = re.search(r"\d+", table_id_raw)
        table_id = str(int(m.group())) if m else table_id_raw
        label = f"Table {table_id}"

    parts = [f"# {label}", "", f"![](../../{image_relpath})", ""]
    if caption:
        parts.append("## Caption")
        parts.append(normalize_ocr_math_text(caption))
    note_texts = [normalize_ocr_math_text(t) for t in table.get("note_texts", []) if t]
    if note_texts:
        parts.append("")
        parts.append("## Notes")
        parts.extend(note_texts)
    if table.get("page"):
        parts.append("")
        parts.append(f"*Page {table['page']}*")
    parts.append("")
    parts.append("---")

    return "\n".join(parts)


def _write_object_markdown(md: str, dst: Path) -> None:
    dst.parent.mkdir(parents=True, exist_ok=True)
    dst.write_text(md.strip() + "\n", encoding="utf-8")


def _find_cached_page_image(page_cache_dir: Path | None, page_num: int) -> Path | None:
    if page_cache_dir is None:
        return None
    for suffix in (".jpg", ".png"):
        candidate = page_cache_dir / f"page_{page_num:03d}{suffix}"
        if candidate.exists():
            return candidate
    return None


def _crop_asset_from_pdf(
    pdf_path: Path | None,
    page_num: int,
    bbox: list[float],
    dst: Path,
    *,
    page_width: int = 0,
    page_height: int = 0,
    page_cache_dir: Path | None = None,
    pdf_doc: Any | None = None,
    pdf_doc_provider: Callable[[], Any | None] | None = None,
    rotation_deg: int = 0,
) -> bool:



    if dst.exists():
        with contextlib.suppress(Exception):
            dst.unlink()

    cached_page_image = _find_cached_page_image(page_cache_dir, page_num)
    if cached_page_image is not None and not rotation_deg:
        try:
            from paperforge.worker.ocr import crop_block_asset
        except ImportError:
            return False
        ok = crop_block_asset(cached_page_image, [int(v) for v in bbox], dst)
        return ok

    created_doc = None
    doc = pdf_doc
    if doc is None and pdf_doc_provider is not None:
        doc = pdf_doc_provider()
        if doc is None:
            return False
    if doc is None:
        if pdf_path is None or not pdf_path.exists():
            return False
        try:
            import fitz
        except ImportError:
            return False
        created_doc = fitz.open(str(pdf_path))
        doc = created_doc

    try:
        if page_width > 0 and page_height > 0 and page_cache_dir is not None and not rotation_deg:
            try:
                from paperforge.worker.ocr import crop_block_asset, render_pdf_page_cached
            except ImportError:
                return False

            try:
                page_image_path = page_cache_dir / f"page_{page_num:03d}.jpg"
                rendered = render_pdf_page_cached(
                    doc,
                    page_num,
                    target_width=page_width,
                    target_height=page_height,
                    destination=page_image_path,
                )
                if not rendered:
                    return False
                ok = crop_block_asset(rendered, [int(v) for v in bbox], dst)
                return ok
            except Exception:
                return False

        try:
            import fitz
        except ImportError:
            return False
        try:
            page = doc[page_num - 1]
            # Clip is in PDF user space — bbox is in OCR coordinates (scaled).
            # Convert to PDF space using page_width/height vs PDF page rect ratio.
            pdf_rect = page.rect
            sx = max(1.0, page_width / pdf_rect.width) if page_width > 0 else 2.0
            sy = max(1.0, page_height / pdf_rect.height) if page_height > 0 else 2.0
            rect = fitz.Rect(bbox[0] / sx, bbox[1] / sy, bbox[2] / sx, bbox[3] / sy)
            # High-resolution zoom for crisp vector rendering (4x vs previous 2x)
            zoom = 4.0
            mat = fitz.Matrix(zoom, zoom)
            pix = page.get_pixmap(matrix=mat, clip=rect)
            dst.parent.mkdir(parents=True, exist_ok=True)
            if rotation_deg:
                # PyMuPDF Pixmap has no .rotate() — render at 4x then rotate via PIL
                # on the high-res pixmap data. Single pass, no double JPEG compression.
                from PIL import Image
                import io
                img = Image.open(io.BytesIO(pix.tobytes("png")))
                img = img.rotate(rotation_deg, expand=True, resample=Image.Resampling.LANCZOS)
                img.save(str(dst))
            else:
                pix.save(str(dst))
            return True
        except Exception:
            return False
    finally:
        if created_doc is not None:
            doc.close()


def extract_and_write_objects(
    pdf_path: Path | None,
    figure_inventory: dict[str, Any],
    table_inventory: dict[str, Any],
    asset_root: Path,
    render_root: Path,
    *,
    page_dimensions_by_page: dict[int, tuple[int, int]] | None = None,
) -> None:
    """Extract figure/table asset crops from PDF and write object markdown."""
    # Validation-first figure matching may retain held figures in inventory,
    # but object emission remains limited to matched figures and unresolved
    # media clusters until figure evidence is sufficient.
    figures_asset_dir = asset_root / "figures"
    tables_asset_dir = asset_root / "tables"
    orphans_asset_dir = asset_root / "orphans"
    figures_render_dir = render_root / "figures"
    tables_render_dir = render_root / "tables"
    page_cache_dir = asset_root.parent / "pages"

    for d in (
        figures_asset_dir,
        tables_asset_dir,
        orphans_asset_dir,
        figures_render_dir,
        tables_render_dir,
        page_cache_dir,
    ):
        d.mkdir(parents=True, exist_ok=True)

    for stale_dir, pattern in (
        (figures_asset_dir, "*.jpg"),
        (tables_asset_dir, "*.jpg"),
        (orphans_asset_dir, "*.jpg"),
        (figures_render_dir, "*.md"),
        (tables_render_dir, "*.md"),
    ):
        for stale in stale_dir.glob(pattern):
            with contextlib.suppress(Exception):
                stale.unlink()

    if page_dimensions_by_page is None:
        page_dimensions_by_page = {}

    def _page_dims(page_num: int) -> tuple[int, int]:
        return page_dimensions_by_page.get(page_num, (0, 0))

    shared_pdf_doc: Any | None = None
    shared_pdf_open_attempted = False

    def _get_shared_pdf_doc() -> Any | None:
        nonlocal shared_pdf_doc, shared_pdf_open_attempted
        if shared_pdf_doc is not None:
            return shared_pdf_doc
        if shared_pdf_open_attempted:
            return None
        if pdf_path is None or not pdf_path.exists():
            shared_pdf_open_attempted = True
            return None
        shared_pdf_open_attempted = True
        try:
            import fitz
            shared_pdf_doc = fitz.open(str(pdf_path))
        except Exception:
            return None
        return shared_pdf_doc

    try:
        # Process matched figures
        for i, match in enumerate(figure_inventory.get("matched_figures", [])):
            fig_id = match.get("figure_id", f"figure_{i + 1:03d}")
            caption_text = match.get("text", "")
            page = match.get("page", 0)
            page_width, page_height = _page_dims(page)
            asset_path_rel = f"assets/figures/{fig_id}.jpg"
            asset_path_abs = figures_asset_dir / f"{fig_id}.jpg"

            rotation_deg = int(match.get("rotation_correction_deg", 0) or 0)

            was_cropped = False
            cluster_bbox = match.get("cluster_bbox")
            if cluster_bbox and all(v > 0 for v in cluster_bbox):
                was_cropped = _crop_asset_from_pdf(
                    pdf_path,
                    page,
                    cluster_bbox,
                    asset_path_abs,
                    page_width=page_width,
                    page_height=page_height,
                    page_cache_dir=page_cache_dir,
                    pdf_doc_provider=_get_shared_pdf_doc,
                    rotation_deg=rotation_deg,
                )
            if not was_cropped:
                for asset_info in match.get("matched_assets", []):
                    bbox = asset_info.get("bbox", [0, 0, 0, 0])
                    if (
                        pdf_path
                        and bbox
                        and all(v > 0 for v in bbox)
                        and _crop_asset_from_pdf(
                            pdf_path,
                            page,
                            bbox,
                            asset_path_abs,
                            page_width=page_width,
                            page_height=page_height,
                            page_cache_dir=page_cache_dir,
                            pdf_doc_provider=_get_shared_pdf_doc,
                            rotation_deg=rotation_deg,
                        )
                    ):
                        was_cropped = True
                        break

            md = render_figure_object_markdown(
                {
                    "figure_id": fig_id,
                    "page": page,
                    "caption": caption_text,
                    "image_relpath": asset_path_rel,
                    "confidence": match.get("confidence", 0.5),
                    "was_cropped": was_cropped,
                }
            )
            _write_object_markdown(md, figures_render_dir / f"{fig_id}.md")

        # Process unresolved figure clusters (multi-panel without reliable legend)
        for i, cluster in enumerate(figure_inventory.get("unresolved_clusters", [])):
            cluster_id = cluster.get("cluster_id") or f"unresolved_cluster_{i + 1:03d}"
            page = cluster.get("page", 0)
            bbox = cluster.get("cluster_bbox", [0, 0, 0, 0])
            page_width, page_height = _page_dims(page)
            asset_path_rel = f"assets/figures/{cluster_id}.jpg"
            asset_path_abs = figures_asset_dir / f"{cluster_id}.jpg"

            if pdf_path and bbox and all(v > 0 for v in bbox):
                _crop_asset_from_pdf(
                    pdf_path,
                    page,
                    bbox,
                    asset_path_abs,
                    page_width=page_width,
                    page_height=page_height,
                    page_cache_dir=page_cache_dir,
                    pdf_doc_provider=_get_shared_pdf_doc,
                )

            md = render_figure_object_markdown(
                {
                    "figure_id": cluster_id,
                    "page": page,
                    "caption": "",
                    "image_relpath": asset_path_rel,
                    "confidence": 0.45,
                }
            )
            _write_object_markdown(md, figures_render_dir / f"{cluster_id}.md")

        # Process unmatched assets as orphans
        orphan_count = 0
        for asset in figure_inventory.get("unmatched_assets", []):
            orphan_count += 1
            orphan_id = f"orphan_{orphan_count:03d}"
            bbox = asset.get("bbox", [0, 0, 0, 0])
            page = asset.get("page", 0)
            page_width, page_height = _page_dims(page)
            asset_path_rel = f"assets/orphans/{orphan_id}.jpg"
            asset_path_abs = orphans_asset_dir / f"{orphan_id}.jpg"

            if pdf_path and bbox and all(v > 0 for v in bbox):
                _crop_asset_from_pdf(
                    pdf_path,
                    page,
                    bbox,
                    asset_path_abs,
                    page_width=page_width,
                    page_height=page_height,
                    page_cache_dir=page_cache_dir,
                    pdf_doc_provider=_get_shared_pdf_doc,
                )

            md = render_figure_object_markdown(
                {
                    "figure_id": orphan_id,
                    "page": page,
                    "caption": "",
                    "image_relpath": asset_path_rel,
                    "confidence": 0.3,
                }
            )
            _write_object_markdown(md, figures_render_dir / f"{orphan_id}.md")

        # Process tables
        for i, table in enumerate(table_inventory.get("tables", [])):
            tbl_id = f"table_{i + 1:03d}"
            caption_text = table.get("caption_text", "")
            page = table.get("page", 0)
            page_width, page_height = _page_dims(page)
            asset_bbox = table.get("asset_bbox", [0, 0, 0, 0])
            asset_path_rel = f"assets/tables/{tbl_id}.jpg"
            asset_path_abs = tables_asset_dir / f"{tbl_id}.jpg"

            was_cropped = False
            # Use render_bbox + render_rotation_deg for rotated tables
            _crop_bbox = table.get("render_bbox") or asset_bbox
            _rot_deg = table.get("render_rotation_deg", 0) or 0
            if table.get("has_asset") and pdf_path and _crop_bbox and all(v > 0 for v in _crop_bbox):
                was_cropped = _crop_asset_from_pdf(
                    pdf_path,
                    page,
                    _crop_bbox,
                    asset_path_abs,
                    page_width=page_width,
                    page_height=page_height,
                    page_cache_dir=page_cache_dir,
                    pdf_doc_provider=_get_shared_pdf_doc,
                    rotation_deg=_rot_deg,
                )

            md = render_table_object_markdown(
                {
                    "table_id": tbl_id,
                    "page": page,
                    "caption": caption_text,
                    "image_relpath": asset_path_rel,
                    "confidence": 0.85 if was_cropped else 0.4,
                    "formal_table_number": table.get("formal_table_number") or table.get("table_number"),
                    "note_texts": table.get("note_texts", []),
                    "note_match_reason": table.get("note_match_reason", ""),
                }
            )
            _write_object_markdown(md, tables_render_dir / f"{tbl_id}.md")

        # Process unmatched table assets as orphans
        for asset in table_inventory.get("unmatched_assets", []):
            orphan_count += 1
            orphan_id = f"orphan_{orphan_count:03d}"
            bbox = asset.get("bbox", [0, 0, 0, 0])
            page = asset.get("page", 0)
            page_width, page_height = _page_dims(page)
            asset_path_rel = f"assets/orphans/{orphan_id}.jpg"
            asset_path_abs = orphans_asset_dir / f"{orphan_id}.jpg"

            if pdf_path and bbox and all(v > 0 for v in bbox):
                _crop_asset_from_pdf(
                    pdf_path,
                    page,
                    bbox,
                    asset_path_abs,
                    page_width=page_width,
                    page_height=page_height,
                    page_cache_dir=page_cache_dir,
                    pdf_doc_provider=_get_shared_pdf_doc,
                )

            md = render_figure_object_markdown(
                {
                    "figure_id": orphan_id,
                    "page": page,
                    "caption": "",
                    "image_relpath": asset_path_rel,
                    "confidence": 0.3,
                }
            )
            _write_object_markdown(md, figures_render_dir / f"{orphan_id}.md")
    finally:
        if shared_pdf_doc is not None:
            with contextlib.suppress(Exception):
                shared_pdf_doc.close()