Blame for backend/media_analysis/figure_analysis.py

"""Resolve selected scientific paper media using the EXISTING paper_figures module.

This adapter is isolated from existing chat and routing. Production chat integration
should happen only after verifying IDs, manifest records and endpoint behavior.
"""
from __future__ import annotations

import os
import re
from typing import Any, Callable


class MediaSelectionError(ValueError):
    pass


def _canonical_label(label: str) -> tuple[str, int] | None:
    # Extracted manifest labels use figure_03 / table_04, while users ask
    # for Figure 3 / Fig. 3 / Table 4. Normalize both without touching manifests.
    match = re.search(r"(?<!\w)(fig(?:ure)?|table)[\s_.-]*(\d+)(?!\d)", label or "", re.I)
    if not match:
        return None
    kind = "table" if match.group(1).lower() == "table" else "figure"
    return kind, int(match.group(2))


def _requested_labels(question: str) -> list[tuple[str, int]]:
    found = []
    for match in re.finditer(r"\b(fig(?:ure)?\.?|table)\s*(\d+)\b", question, re.I):
        label = ("table" if match.group(1).lower() == "table" else "figure", int(match.group(2)))
        if label not in found:
            found.append(label)
    return found


def prepare_media_evidence(
    question: str,
    selected_papers: list[dict[str, Any]],
    *,
    figures_lookup: Callable[..., dict | None] | None = None,
    asset_resolver: Callable[[str, str], Any] | None = None,
    max_images: int = 4,
    media_image: str | None = None,
) -> list[dict]:
    """Explicit label-matching POC; never retrieves unselected papers.

    selected_papers must contain paper_id and at least one of paper_key/txt_file/
    pdf_file/metadata_row_index. Optional paper_title, doi, summary are carried
    without performing new retrieval. Ambiguous figure labels across selected
    papers are deliberately rejected (identify the paper or pass one paper).
    """
    if figures_lookup is None or asset_resolver is None:
        from metadata_search.paper_figures import figures_payload_for_paper, resolve_figure_asset
        figures_lookup = figures_lookup or figures_payload_for_paper
        asset_resolver = asset_resolver or resolve_figure_asset
    if not selected_papers:
        raise MediaSelectionError("Select at least one paper")
    if max_images < 1 or max_images > 4:
        raise ValueError("max_images must be 1-4 for the POC")
    labels = _requested_labels(question)
    if not labels:
        raise MediaSelectionError("POC requires an explicit Figure N or Table N reference")
    if len(selected_papers) > 4:
        raise MediaSelectionError("POC supports up to four selected papers")

    matched = []
    for paper in selected_papers:
        paper_id = str(paper.get("paper_id") or paper.get("document_id") or "").strip()
        if not paper_id:
            raise MediaSelectionError("Each selected paper needs a stable paper_id/document_id")
        lookup_args = {key: paper[key] for key in ("txt_file", "pdf_file", "paper_key", "metadata_row_index")
                       if paper.get(key) is not None and paper.get(key) != ""}
        if not lookup_args:
            raise MediaSelectionError(f"Paper {paper_id} missing a media lookup locator")
        payload = figures_lookup(**lookup_args)
        if not payload:
            continue
        folder = payload.get("paper_key")
        if not folder:
            continue
        for entry in payload.get("items") or []:
            actual = _canonical_label(entry.get("label", ""))
            if actual not in labels:
                continue
            # Reject duplicated figure labels inside one paper rather than guessing.
            image = entry.get("image")
            if media_image is not None and image != media_image:
                continue
            if not image:
                continue
            image_path = asset_resolver(folder, image)
            if image_path is None:
                raise MediaSelectionError(f"Image missing for paper {paper_id}, {entry.get('label')}")
            matched.append({
                "paper_id": paper_id,
                "paper_title": paper.get("paper_title") or paper.get("title", ""),
                "doi": paper.get("doi", ""),
                "summary": paper.get("summary", ""),
                "paper_key": folder,
                "label": entry.get("label", ""),
                "panel": "",  # Full crop first; panel-specific cropping is not available.
                "page": entry.get("page"),
                "caption": entry.get("caption", ""),
                "image_url": entry.get("image_url", ""),
                "image_path": str(image_path),
                "kind": entry.get("kind", "figure"),
            })
    if not matched:
        raise MediaSelectionError("Requested image not found for this paper and figure/table label" if media_image else "Requested media not found in the selected papers")
    # Do not conflate 'Figure 3' across multiple papers without a paper qualifier.
    if len(selected_papers) > 1 and len(labels) == 1 and len(matched) > 1:
        # Explicit cross-paper comparisons are allowed; otherwise require disambiguation.
        if not re.search(r"\b(compare|across|between|both|all|different studies|selected papers)\b", question, re.I):
            raise MediaSelectionError("Figure/table label occurs in several selected papers; specify paper or ask an explicit cross-paper comparison")
    if len(matched) > max_images:
        raise MediaSelectionError(f"Matched {len(matched)} images; POC limit is {max_images}. Narrow the request")
    keys = [(x["paper_id"], _canonical_label(x["label"])) for x in matched]
    if len(set(keys)) != len(keys):
        raise MediaSelectionError("Multiple images share this figure/table number. Click Analyze & CSV under the intended image in the gallery.")
    return matched


def analyze_media_question(question: str, selected_papers: list[dict], *, enabled: bool | None = None,
                           figures_lookup=None, asset_resolver=None, vision_client=None, provider: str = "azure", model: str | None = None, max_output_tokens: int | None = None,
                           evidence_mode: str = "vision", export_csv: bool = False, analysis_mode: str = "auto", media_image: str | None = None) -> dict:
    """Explicit opt-in entry point. Never silently substitutes a text summary answer."""
    if enabled is None:
        enabled = os.getenv("WATERSCOPE_ENABLE_MEDIA_ANALYSIS", "0").lower() in ("1", "true", "yes")
    if not enabled:
        return {"status": "disabled", "answer": None, "vision_inspected": False}
    items = prepare_media_evidence(question, selected_papers,
                                   figures_lookup=figures_lookup, asset_resolver=asset_resolver, media_image=media_image)
    if vision_client is None:
        from .vision_client import analyze_images
        vision_client = analyze_images
    if analysis_mode not in ("auto", "extract_all", "interpret"):
        raise ValueError("analysis_mode must be auto, extract_all, or interpret")
    if analysis_mode == "extract_all":
        question = question + ("\n[Task: extract all legitimately readable values from the selected media; "
                               "do not invent chart points or force numerical estimates where unreadable.]")
    elif analysis_mode == "interpret":
        question = question + "\n[Task: prioritize interpretation; extract only data relevant to the question.]"
    if evidence_mode not in ("vision", "json", "combined"):
        raise ValueError("evidence_mode must be vision, json or combined")
    # Figure images usually have no positional JSON; combined mode automatically
    # uses the image and caption without requiring a fictitious figure JSON.
    if evidence_mode != "vision":
        from .table_data import load_table_evidence
        if asset_resolver is None:
            from metadata_search.paper_figures import resolve_figure_asset
            asset_resolver = resolve_figure_asset
        for item in items:
            if item.get("kind") == "table":
                item["table_data"] = load_table_evidence(item["paper_key"], item["label"], asset_resolver)
        if evidence_mode == "json" and any(not x.get("table_data") for x in items):
            raise MediaSelectionError("JSON-only mode requires extracted table JSON for every requested item")
    result = vision_client(question, items, provider=provider, model=model,
                           max_output_tokens=max_output_tokens, evidence_mode=evidence_mode,
                           structured_output=export_csv)
    if export_csv:
        from .data_export import decode_structured_answer, csv_attachment, normalize_records
        if result.get("response_status") == "incomplete" or result.get("incomplete_reason"):
            answer, records, warning = "", [], "Model output incomplete; CSV was not generated"
        else:
            answer, records, warning = decode_structured_answer(result.get("answer", ""))
        result["answer"] = answer if not warning else (
            "Structured extraction did not complete; no downloadable CSV was generated. "
            "Please retry or narrow the request."
        )
        records, normalization_warnings = normalize_records(records)
        result["structured_data"] = records
        result["normalization_warnings"] = normalization_warnings
        sources = result.get("sources") or []
        if len(sources) > 1 and records:
            names = {key.lower() for rec in records for key in rec}
            if not any(name in names for name in ("source", "source_id", "paper_id", "paper_title", "source_media_label", "media_label", "source_label")):
                warning = warning or "Multiple sources require a source identifier column; CSV not generated"
                records = []
                result["structured_data"] = []
        result["csv_attachment"] = csv_attachment(records, sources)
        if warning:
            result["export_warning"] = warning
    return {"status": "ok", "vision_inspected": evidence_mode != "json", "evidence_mode": evidence_mode, "analysis_mode": analysis_mode, **result}