noteman-slicer/tests/test_render.py

"""Runnable check for the render pipeline and bundle export."""

from __future__ import annotations

import json
import sys
import zipfile
from pathlib import Path

import cv2
import numpy as np
import pymupdf

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from noteman_slicer import bundle  # noqa: E402
from noteman_slicer.detect import detect_page  # noqa: E402
from noteman_slicer.pdf import open_source, page_raster  # noqa: E402
from noteman_slicer.project import Cut, Project, default_path  # noqa: E402
from noteman_slicer.render import (  # noqa: E402
    ALPHA_LEVELS,
    apply_levels,
    encode,
    pad_right,
    render_slices,
    scale_song,
)

W, H = 1200, 1600
GAP = 15


def _system(page: np.ndarray, top: int, right: int) -> None:
    """A bracket plus two staves, with a lyric line under each."""
    page[top : top + 200, 100:104] = 0
    for staff in (top, top + 140):
        for i in range(5):
            page[staff + i * GAP : staff + i * GAP + 2, 110:right] = 0
        page[staff + 90 : staff + 105, 200 : right - 100] = 0


def _scan_pdf(path: Path) -> None:
    art = np.full((H, W), 255, np.uint8)
    art[40:60, 400:800] = 0  # title, far from any system
    _system(art, 300, 1100)
    _system(art, 800, 900)  # narrower: exercises the right pad
    pix = pymupdf.Pixmap(pymupdf.csGRAY, W, H, bytearray(art.tobytes()), False)
    doc = pymupdf.open()
    page = doc.new_page(width=595, height=842)
    page.insert_image(page.rect, pixmap=pix)
    doc.save(path)


def main() -> int:
    tmp = Path(__file__).with_name("_tmp")
    tmp.mkdir(exist_ok=True)
    pdf = tmp / "scan.pdf"
    _scan_pdf(pdf)

    source = open_source(pdf)
    gray = page_raster(source, 0)
    project = Project.from_detection(pdf, [detect_page(gray)], [gray.shape[0]])

    slices = render_slices(project, source)
    assert len(slices) == 2, f"expected 2 kept slices, got {len(slices)}"

    # The title is far from any bracket, so it is not in a kept slice: both
    # slices must be shorter than the gap between the systems.
    assert all(s.gray.shape[0] < 400 for s in slices), [s.gray.shape for s in slices]

    # System 2 is drawn narrower, so before padding the widths differ.
    assert slices[0].width != slices[1].width, "the fixture should differ in width"

    scaled = scale_song(slices, cap=4000)  # a cap far above the fixture
    assert all(abs(a.shape[1] - b.width) <= 2 for a, b in zip(scaled, slices)), (
        "never upscale: a song narrower than the cap must be left alone"
    )

    padded = pad_right(scale_song(slices))
    assert len({p.shape[1] for p in padded}) == 1, "slices must share one width"
    assert max(p.shape[1] for p in padded) <= 1920

    rgba = cv2.imdecode(np.frombuffer(encode(padded[0]), np.uint8), cv2.IMREAD_UNCHANGED)
    assert rgba.shape[2] == 4
    assert rgba[:, :, :3].max() == 0, "ink must be pure black"
    assert rgba[:, :, 3].max() == 255, "full ink must be fully opaque"
    assert rgba[:, :, 3].min() == 0, "paper must be fully transparent"
    assert len(np.unique(rgba[:, :, 3])) <= ALPHA_LEVELS

    # Levels: a white point below the paper value wipes the paper out entirely.
    faint = np.full((10, 10), 200, np.uint8)
    assert apply_levels(faint, 0, 180).max() == 255

    # The Engel case: a section label printed in the left margin at a height
    # that belongs to the *next* system. A straight cut cannot separate it from
    # the previous system's lyrics; a stepped one can.
    label_top, label_bottom = 620, 680
    labelled = tmp / "labelled.pdf"
    art = np.full((H, W), 255, np.uint8)
    _system(art, 300, 1100)
    _system(art, 800, 900)
    art[label_top:label_bottom, 120:300] = 0  # the label
    art[label_top:label_bottom, 500:1000] = 0  # system 1's trailing lyrics, same rows
    pix = pymupdf.Pixmap(pymupdf.csGRAY, W, H, bytearray(art.tobytes()), False)
    doc = pymupdf.open()
    doc.new_page(width=595, height=842).insert_image(pymupdf.Rect(0, 0, 595, 842), pixmap=pix)
    doc.save(labelled)

    src2 = open_source(labelled)
    g2 = page_raster(src2, 0)
    proj2 = Project.from_detection(labelled, [detect_page(g2)], [g2.shape[0]])
    page = proj2.pages[0]
    scale = g2.shape[0] / H

    def ink(images: list) -> list[int]:
        """Ink in the left margin of each slice — where the label sits."""
        return [int((i.gray[:, : int(i.width * 0.3)] < 128).sum()) for i in images]

    # Straight cut through the middle of that band: the label goes with
    # whichever side the line falls on, and cannot be separated.
    band_mid = (label_top + label_bottom) / 2 * scale / g2.shape[0]
    page.cuts[1] = Cut.straight(band_mid)
    straight_ink = ink(render_slices(proj2, src2))

    # Stepped: above the label on the left, below the lyrics on the right.
    above = (label_top - 10) * scale / g2.shape[0]
    below = (label_bottom + 10) * scale / g2.shape[0]
    page.cuts[1] = Cut([(0.0, above), (0.35, above), (0.35, below), (1.0, below)])
    stepped_ink = ink(render_slices(proj2, src2))

    # The straight cut splits the label down the middle; the stepped cut gives
    # all of it to the lower slice and none to the upper.
    assert stepped_ink[1] > straight_ink[1], (
        f"the label must move into the lower slice: {straight_ink} → {stepped_ink}"
    )
    assert stepped_ink[0] < straight_ink[0], (
        f"and out of the upper one: {straight_ink} → {stepped_ink}"
    )
    src2.close()
    labelled.unlink()

    # Bundle. A title is required; everything else is optional.
    try:
        bundle.write(project, source, tmp / "untitled.zip")
    except ValueError as error:
        assert "title" in str(error)
    else:
        raise AssertionError("export without a title should be refused")

    project.metadata["title"] = "Test song"
    out = bundle.write(project, source, tmp / "song.zip")
    with zipfile.ZipFile(out) as zf:
        names = zf.namelist()
        assert "song.json" in names and "original.pdf" in names, names
        meta = json.loads(zf.read("song.json"))
        assert meta["v"] == 1
        files = [s["file"] for s in meta["slices"]]
        assert files == ["001.webp", "002.webp"], files
        assert all(f in names for f in files)

    source.close()
    # Exporting marks the project spent, which writes the project file.
    for f in (pdf, out, default_path(pdf)):
        f.unlink(missing_ok=True)
    tmp.rmdir()
    print("ok")
    return 0


if __name__ == "__main__":
    sys.exit(main())