formula_img2latex/tests/test_pipeline.py

"""scan → (candidates) → verify → apply → revert on a synthetic epub2md output, without any model."""

from __future__ import annotations

from pathlib import Path

from conftest import needs_latex, needs_pandoc
from formula_img2latex import pipeline
from formula_img2latex.cli import main
from formula_img2latex.engines.vlm import parse_answer, prepare_image
from formula_img2latex.sources.base import scan_markdown
from formula_img2latex.store import Workspace


def _inject(ws: Workspace, by_asset: dict[str, str]) -> None:
    """Pretend an engine answered (tests must not need a model server).

    The mechanics are tested here, not the calibrated thresholds (test_core.test_decision_rules),
    so the contrast threshold is set low enough for the synthetic images.
    """
    ws.settings["thresholds"] = {"min_contrast": 0.5}
    ws.settings["decide"] = "pixel"
    ws.save_meta()
    for f in ws.all():
        for name, tex in by_asset.items():
            if f.asset.endswith(name):
                f.candidates.append({"engine": "test", "model": "fixture", "tex": tex, "is_formula": True})
                ws.save(f)


def test_scan_markdown_skips_code_and_detects_display(epub2md_dir: Path) -> None:
    refs = scan_markdown(epub2md_dir / "chapters" / "001-kapitel.md")
    assert [(r.target.rsplit("/", 1)[1], r.display) for r in refs] == [
        ("0001-inline.jpg", False), ("0002-display.jpg", True), ("0003-figure.png", True)]


def test_scan_classifies(epub2md_dir: Path) -> None:
    source, ws = pipeline.open_all(epub2md_dir, None, create=True)
    pipeline.scan(source, ws, quiet=True)
    labels = {f.asset.rsplit("/", 1)[1]: (f.classify["label"], f.known_formula) for f in ws.all()}
    assert labels["0001-inline.jpg"] == ("formula", False)
    assert labels["0002-display.jpg"] == ("formula", True)       # epub2md warning
    assert labels["0003-figure.png"][0] == "figure"
    assert ws.root == epub2md_dir.parent / "buch.formulas"


@needs_latex
@needs_pandoc
def test_verify_apply_revert_roundtrip(epub2md_dir: Path) -> None:
    source, ws = pipeline.open_all(epub2md_dir, None, create=True)
    pipeline.scan(source, ws, quiet=True)
    _inject(ws, {"0001-inline.jpg": r"L = \{ - \sqrt{2} , \sqrt{2} \}",          # comma: wrong, repaired to ;
                 "0002-display.jpg": r"x_{2,k}=\arcsin\left(-\frac{3}{5}\right)+2k\pi"})
    counts = pipeline.verify_all(source, ws, preset="cm", workers=2, quiet=True)
    st = {f.asset.rsplit("/", 1)[1]: f for f in ws.all()}
    inline = st["0001-inline.jpg"]
    assert inline.verification["best"]["tex"].replace(" ", "") == r"L=\{-\sqrt{2};\sqrt{2}\}"
    assert inline.status == "auto_accepted", inline.reason
    assert sum(counts.values()) == 2
    md = epub2md_dir / "chapters" / "001-kapitel.md"
    before = md.read_text(encoding="utf-8")

    pipeline.apply(source, ws, quiet=True)
    after = md.read_text(encoding="utf-8")
    assert "lautet $L = \\{ - \\sqrt{2} ; \\sqrt{2} \\}$." in after or "lautet $L=\\{-\\sqrt{2};\\sqrt{2}\\}$." in after
    assert "```\n![kein](../media/images/0001-inline.jpg)\n```" in after      # code block untouched
    assert "0003-figure.png" in after
    pipeline.apply(source, ws, quiet=True)                                        # idempotent
    assert md.read_text(encoding="utf-8") == after

    pipeline.revert(source, ws, quiet=True)
    assert md.read_text(encoding="utf-8") == before


@needs_latex
def test_revert_refuses_hand_edited_files(epub2md_dir: Path) -> None:
    source, ws = pipeline.open_all(epub2md_dir, None, create=True)
    pipeline.scan(source, ws, quiet=True)
    _inject(ws, {"0001-inline.jpg": r"L=\{-\sqrt{2};\sqrt{2}\}"})
    pipeline.verify_all(source, ws, preset="cm", workers=1, quiet=True)
    pipeline.apply(source, ws, quiet=True)
    md = epub2md_dir / "chapters" / "001-kapitel.md"
    md.write_text(md.read_text(encoding="utf-8") + "\nHandnotiz\n", encoding="utf-8")
    assert pipeline.revert(source, ws, quiet=True)["skipped"] == 1
    assert "Handnotiz" in md.read_text(encoding="utf-8")


def test_cli_errors(tmp_path: Path, capsys) -> None:
    assert main(["status", str(tmp_path / "fehlt")]) == 1
    (tmp_path / "leer").mkdir()
    assert main(["scan", str(tmp_path / "leer")]) == 2
    assert "weder eine epub2md-Ausgabe" in capsys.readouterr().err


def test_engine_answer_parsing() -> None:
    assert parse_answer('{"is_formula": true, "latex": "x^2"}') == ("x^2", True)
    assert parse_answer('```json\n{"is_formula": true, "latex": "\\frac{a}{b}"}\n```') == ("\\frac{a}{b}", True)
    assert parse_answer('{"is_formula": true, "latex": "\\sqrt{2}"}') == ("\\sqrt{2}", True)   # unescaped backslash
    assert parse_answer('{"is_formula": false, "latex": ""}') == (None, False)
    assert parse_answer("x_1 + x_2") == ("x_1 + x_2", None)


def test_prepare_image_enlarges_small_formulas(epub2md_dir: Path) -> None:
    import io

    from PIL import Image
    data = (epub2md_dir / "media" / "images" / "0001-inline.jpg").read_bytes()
    with Image.open(io.BytesIO(prepare_image(data, "x.jpg"))) as im:
        assert im.height >= 96 and im.mode == "L"


# ------------------------------------------------------------------------------- consensus

def _reading(engine: str, model: str, tex: str | None, p_mean: float = 0.99, p_min: float = 0.8,
             is_formula: bool = True) -> dict:
    return {"engine": engine, "model": model, "tex": tex, "is_formula": is_formula, "p_mean": p_mean, "p_min": p_min}


@needs_latex
@needs_pandoc
def test_consensus_mode(epub2md_dir: Path) -> None:
    source, ws = pipeline.open_all(epub2md_dir, None, create=True)
    pipeline.scan(source, ws, quiet=True)
    by = {f.asset.rsplit("/", 1)[1]: f for f in ws.all()}
    inline, display = by["0001-inline.jpg"], by["0002-display.jpg"]
    inline.candidates = [_reading("local", "A", r"L=\{-\sqrt{2};\sqrt{2}\}"),
                         _reading("local2", "B", r"L = \{ - \sqrt2 ; \sqrt2 \}")]          # same, other spelling
    display.candidates = [_reading("local", "A", r"x_{2,k}=\arcsin\left(-\frac{3}{5}\right)+2k\pi"),
                          _reading("local2", "B", r"x_{3,k}=\arcsin\left(-\frac{3}{5}\right)+2k\pi")]  # differ
    ws.save(inline)
    ws.save(display)
    pipeline.verify_all(source, ws, preset="cm", workers=2, quiet=True)
    st = {f.asset.rsplit("/", 1)[1]: f for f in ws.all()}
    assert pipeline.needs_third(st["0001-inline.jpg"]), st["0001-inline.jpg"].reason    # A = B, C missing
    assert st["0002-display.jpg"].status == "needs_review"
    assert "uneinig" in st["0002-display.jpg"].reason

    # the third reading confirms A = B -> accepted; for A != B a third reading does not decide 2:1
    i, d = st["0001-inline.jpg"], st["0002-display.jpg"]
    i.candidates.append(_reading("cloud", "C", r"L=\left\{-\sqrt{2};\sqrt{2}\right\}"))
    d.candidates.append(_reading("cloud", "C", r"x_{2,k} = \arcsin(-\frac35) + 2k\pi"))
    ws.save(i)
    ws.save(d)
    pipeline.verify_all(source, ws, preset="cm", workers=2, quiet=True)
    i, d = ws.load(i.id), ws.load(d.id)
    assert i.status == "auto_accepted" and "dritte Lesart bestätigt" in i.reason, i.reason
    assert d.status == "needs_review"

    # a disagreeing third reading blocks the acceptance
    i.candidates[-1] = _reading("cloud", "C", r"L=\{-\sqrt{2};\sqrt{3}\}")
    i.status = "pending"
    ws.save(i)
    pipeline.verify_all(source, ws, preset="cm", workers=2, quiet=True)
    i = ws.load(i.id)
    assert i.status == "needs_review" and "weicht ab" in i.reason, i.reason


def test_consensus_rules() -> None:
    from formula_img2latex.decide import consensus_thresholds, decide_consensus
    th = consensus_thresholds({"min_p_mean": 0.9, "min_p_min": 0.3, "require_third": 0})
    a, b = _reading("local", "A", "x^2"), _reading("local2", "B", "x^{2}")
    veto = {"ok": True, "metrics": {"ncc": 0.9, "px_per_pt": 1.4}}
    assert decide_consensus([a, b], [(0, 1, "normalform")], veto, True, 1.4, th)[0] == "auto_accepted"
    assert decide_consensus([a, b], [], veto, True, 1.4, th)[0] == "needs_review"                       # disagree
    low = _reading("local2", "B", "x^{2}", p_mean=0.7)
    assert decide_consensus([a, low], [(0, 1, "normalform")], veto, True, 1.4, th)[0] == "needs_review"  # unsure
    bad = {"ok": True, "metrics": {"ncc": 0.5, "px_per_pt": 1.4}}
    assert decide_consensus([a, b], [(0, 1, "normalform")], bad, True, 1.4, th)[0] == "needs_review"     # veto
    big = {"ok": True, "metrics": {"ncc": 0.9, "px_per_pt": 1.8}}
    sized = {**th, "veto_scale_dev": 0.15}
    assert decide_consensus([a, b], [(0, 1, "normalform")], big, True, 1.4, sized)[0] == "needs_review"  # size check on
    assert decide_consensus([a, b], [(0, 1, "normalform")], veto, False, 1.4, th)[0] == "needs_review"   # pandoc
    nf = [_reading("local", "A", None, is_formula=False), _reading("local2", "B", None, is_formula=False)]
    assert decide_consensus(nf, [], None, True, 1.4, th)[0] == "not_formula"
    assert decide_consensus([a], [], veto, True, 1.4, th)[0] == "needs_review"                          # one model
    strict = consensus_thresholds()                                                     # default: third must agree
    c = _reading("cloud", "C", "x^2")
    assert decide_consensus([a, b], [(0, 1, "normalform")], veto, True, 1.4, strict)[0] == "needs_review"
    assert decide_consensus([a, b, c], [(0, 1, "normalform"), (0, 2, "normalform"), (1, 2, "normalform")],
                            veto, True, 1.4, strict)[0] == "auto_accepted"
    assert decide_consensus([a, b, c], [(0, 2, "normalform")], veto, True, 1.4, strict)[0] == "needs_review"  # no 2:1


def test_canonical_form() -> None:
    from formula_img2latex.consensus import canonical
    same = [(r"x^{2}+y^{2}", "x^2 + y^2"), (r"\vec{b}", r"\overrightarrow{b}"), (r"\left(a\right)", "(a)"),
            (r"\mathrm{d}x", "dx"), (r"\frac{1}{2}", r"\dfrac12"), (r"{\mathrm{e}}^{x}", "e^x"), (r"x\le1", r"x \leq 1"),
            (r"\mathbf{v}", r"\boldsymbol{v}")]
    differ = [(r"x_{2,k}", r"x_{2;k}"), ("x^{2}", "x_{2}"), (r"\mathbb{R}", "R"), ("a_{n+1}", "a_n+1"),
              (r"\frac{a}{b+c}", r"\frac{a}{b}+c"), (r"x \in t", r"x \int"), (r"\sqrt[3]{x}", r"\sqrt{x}")]
    assert all(canonical(a) == canonical(b) for a, b in same)
    assert all(canonical(a) != canonical(b) for a, b in differ)


def test_logprob_confidence() -> None:
    import math

    from formula_img2latex.engines.vlm import latex_token_probs, parse_confidence
    content = '{"is_formula": true, "latex": "x^{2}", "confidence": 97}'
    toks = ['{"', "is", "_formula", '": ', "true", ', "', "latex", '": "', "x", "^{", "2", "}", '", "', "confidence",
            '": ', "97", "}"]
    items = [{"token": t, "logprob": math.log(0.5 if t == "2" else 0.99)} for t in toks]
    p_min, p_mean = latex_token_probs(content, items)
    assert p_min == 0.5 and 0.8 < p_mean < 0.9          # only the four LaTeX tokens count
    assert parse_confidence(content) == 97


def test_equation_number_only_from_chosen_reading() -> None:
    from formula_img2latex.store import Formula
    f = Formula(id="f-x-d", sha256="x", asset="a.png", display=True)
    f.candidates = [{"engine": "local", "model": "old", "tex": r"a=b \tag{3.7}"}]    # an old, hallucinated number
    assert pipeline._with_tag(f, "a=b") == "a=b"
    assert pipeline._with_tag(f, r"a=b \tag{2.1}") == r"a=b\qquad\text{(2.1)}"
    f.display = False
    assert pipeline._with_tag(f, r"a=b \tag{2.1}") == "a=b"


def test_german_decimal_comma(epub2md_dir: Path) -> None:
    import json
    source, _ws = pipeline.open_all(epub2md_dir, None, create=True)
    assert source.format_math("3,6 + (1,2)", False) == "$3,6 + (1,2)$"          # no language: unchanged
    rep = json.loads((epub2md_dir / "conversion-report.json").read_text())
    rep["metadata"] = {"language": "de"}
    (epub2md_dir / "conversion-report.json").write_text(json.dumps(rep))
    source, _ws = pipeline.open_all(epub2md_dir, None)
    assert source.format_math(r"\frac{72}{3,6} = 20, x", True) == r"$$\frac{72}{3{,}6} = 20, x$$"


def test_verify_skips_formulas_without_current_reading(epub2md_dir: Path) -> None:
    source, ws = pipeline.open_all(epub2md_dir, None, create=True)
    pipeline.scan(source, ws, quiet=True)
    f = next(f for f in ws.all() if f.asset.endswith("0001-inline.jpg"))
    f.candidates = [{"engine": "local", "model": "A", "tex": "x", "is_formula": True, "prompt": 1},   # old reading
                    {"engine": "local", "model": "A", "tex": None, "is_formula": False, "prompt": 2}]  # current: none
    ws.save(f)
    pipeline.verify_all(source, ws, preset="cm", workers=1, quiet=True)    # must not crash