sras-viewer/tests/test_batch_export.py

"""Batch Export View as Images: does the exported PNG actually match what
the live view would show, and does the batch dispatch (menu action ->
worker -> per-file render) behave like Batch Compute's proven pattern?

sras_render.export_view_image is tested directly (no Qt) for the plumbing
that decides *what* gets rendered -- pad_factor derived per file, masked-
pixel NaN fill, per-file auto-scale, velocity scaling -- via a
draw_view_image spy rather than pixel-diffing PNGs, the same "spy on the
seam, don't inspect the rendered artifact" approach test_highlight_masked_
pixels (tests/test_gui.py) uses for the live canvas.

The GUI-dispatch half drives SrasViewerWindow._on_batch_export_images()
end-to-end with patched file dialogs, the same shape as
test_stored_cache.py's test_viewer_batch_row_average_dispatch.
"""

from pathlib import Path
from unittest.mock import patch

import numpy as np
import pytest
from PyQt6.QtCore import QEventLoop, QTimer
from PyQt6.QtWidgets import QApplication

import sras_render
from sras_compute import compute_rf_image, dc_image_mv
from sras_format import CH1_IDX, CH3_IDX, CH4_IDX, CH_NAMES, SrasFile
from sras_render import export_view_image
from sras_viewer import SrasViewerWindow, VELOCITY_MODE_IDX
from sras_workers import BatchExportImagesWorker
import tools.make_test_sras as gen

_THRESHOLD_MV = 50.0


def pump(ms: int = 200):
    loop = QEventLoop()
    QTimer.singleShot(ms, loop.quit)
    loop.exec()


def wait_until(pred, timeout_ms: int = 20000, step: int = 100) -> bool:
    waited = 0
    while waited < timeout_ms:
        if pred():
            return True
        pump(step)
        waited += step
    return pred()


# ---------------------------------------------------------------------------
# sras_render.export_view_image -- pure function, no Qt
# ---------------------------------------------------------------------------

_DEFAULT_KW = dict(
    is_fft_mode=False, is_velocity=False, dc_threshold_mv=_THRESHOLD_MV,
    apply_bg_sub=True, pad_factor=1, min_freq_mhz=0.0, grating_um=1.0,
    cmap="viridis", auto_scale=True, vmin=0.0, vmax=1.0,
    highlight_masked=False, mode_str="DC", colorbar_label="mV",
)


def _kw(**overrides):
    kw = dict(_DEFAULT_KW)
    kw.update(overrides)
    return kw


def _spy_draw(monkeypatch):
    """Patches sras_render.draw_view_image to record the image array and
    vmin/vmax/bad_color it was called with, then delegates to the real
    implementation so the PNG is still written -- lets a test check *what*
    export_view_image computed without depending on rendered PNG pixels."""
    orig = sras_render.draw_view_image
    captured = {}

    def spy(ax, fig, img, extent, cmap, vmin, vmax, xlabel, ylabel, title,
            colorbar_label="", cb_ticks=None, norm=None, bad_color=None):
        captured["img"] = np.array(img, copy=True)
        captured["vmin"] = vmin
        captured["vmax"] = vmax
        captured["bad_color"] = bad_color
        return orig(ax, fig, img, extent, cmap, vmin, vmax, xlabel, ylabel,
                   title, colorbar_label, cb_ticks, norm, bad_color)

    monkeypatch.setattr(sras_render, "draw_view_image", spy)
    return captured


def test_export_dc_channel_writes_png(tmp_path):
    path = tmp_path / "dc.sras"
    gen.write(path, n_angles=2, seed=1, samples_per_frame=64)
    out_dir = tmp_path / "out"
    out_dir.mkdir()

    err, out_name = export_view_image(
        str(path), out_dir=str(out_dir), angle_idx=0, ch_idx=CH4_IDX, **_kw())
    assert err == ""
    assert out_name == f"dc_angle0_{CH_NAMES[CH4_IDX]}.png"
    out_path = out_dir / out_name
    assert out_path.exists() and out_path.stat().st_size > 0
    assert out_path.read_bytes()[:8] == b"\x89PNG\r\n\x1a\n", "not a valid PNG"


def test_export_out_of_range_angle(tmp_path):
    path = tmp_path / "short.sras"
    gen.write(path, n_angles=2, seed=2, samples_per_frame=64)
    out_dir = tmp_path / "out"
    out_dir.mkdir()

    err, out_name = export_view_image(
        str(path), out_dir=str(out_dir), angle_idx=5, ch_idx=CH4_IDX, **_kw())
    assert err != "" and "2" in err, "error should mention the file's actual angle count"
    assert out_name == ""
    assert list(out_dir.iterdir()) == [], "no file written for a failed export"


def test_export_fft_mode_matches_compute_rf_image(tmp_path, monkeypatch):
    path = tmp_path / "fft.sras"
    gen.write(path, n_angles=1, seed=3, samples_per_frame=128)
    out_dir = tmp_path / "out"
    out_dir.mkdir()

    captured = _spy_draw(monkeypatch)
    err, _ = export_view_image(
        str(path), out_dir=str(out_dir), angle_idx=0, ch_idx=CH1_IDX,
        **_kw(is_fft_mode=True))
    assert err == ""

    sras = SrasFile(str(path))
    expected = compute_rf_image(sras, 0, dc_threshold_mv=_THRESHOLD_MV,
                                apply_bg_sub=True, n_fft=None, min_freq_mhz=0.0)
    assert np.array_equal(captured["img"], expected)


def test_export_velocity_scales_frequency(tmp_path, monkeypatch):
    path = tmp_path / "vel.sras"
    gen.write(path, n_angles=1, seed=4, samples_per_frame=128)
    out_dir = tmp_path / "out"
    out_dir.mkdir()

    captured = _spy_draw(monkeypatch)
    grating_um = 3.5
    err, _ = export_view_image(
        str(path), out_dir=str(out_dir), angle_idx=0, ch_idx=VELOCITY_MODE_IDX,
        **_kw(is_fft_mode=True, is_velocity=True, grating_um=grating_um))
    assert err == ""

    sras = SrasFile(str(path))
    freq = compute_rf_image(sras, 0, dc_threshold_mv=_THRESHOLD_MV,
                            apply_bg_sub=True, n_fft=None, min_freq_mhz=0.0)
    assert np.array_equal(captured["img"], freq * grating_um)


def test_pad_factor_uses_each_files_own_samples_per_frame(tmp_path, monkeypatch):
    """n_fft must be derived per file from that file's own samples_per_frame,
    never a value carried over from whichever file the caller had open --
    otherwise every file but one in a batch gets silently mis-padded."""
    orig_draw = sras_render.draw_view_image
    out_dir = tmp_path / "out"
    out_dir.mkdir()

    for i, spf in enumerate((64, 256)):
        path = tmp_path / f"pad_{spf}.sras"
        gen.write(path, n_angles=1, seed=5 + i, samples_per_frame=spf)

        captured = {}

        def spy(ax, fig, img, *a, __c=captured, **kw):
            __c["img"] = np.array(img, copy=True)
            return orig_draw(ax, fig, img, *a, **kw)

        monkeypatch.setattr(sras_render, "draw_view_image", spy)

        err, _ = export_view_image(
            str(path), out_dir=str(out_dir), angle_idx=0, ch_idx=CH1_IDX,
            **_kw(is_fft_mode=True, pad_factor=4))
        assert err == ""

        sras = SrasFile(str(path))
        expected = compute_rf_image(sras, 0, dc_threshold_mv=_THRESHOLD_MV,
                                    apply_bg_sub=True, n_fft=spf * 4,
                                    min_freq_mhz=0.0)
        assert np.array_equal(captured["img"], expected), \
            f"samples_per_frame={spf}: n_fft must use this file's own value"


def test_highlight_masked_sets_nan_and_bad_color(tmp_path, monkeypatch):
    path = tmp_path / "mask.sras"
    gen.write(path, n_angles=1, seed=7, samples_per_frame=128)
    out_dir = tmp_path / "out"
    out_dir.mkdir()

    sras = SrasFile(str(path))
    dc4 = dc_image_mv(sras, 0, CH4_IDX)
    threshold = float(np.median(dc4))
    expect_masked = dc4 < threshold
    assert expect_masked.any() and not expect_masked.all(), \
        "fixture threshold should mask some but not all pixels"

    captured = _spy_draw(monkeypatch)
    err, _ = export_view_image(
        str(path), out_dir=str(out_dir), angle_idx=0, ch_idx=CH1_IDX,
        **_kw(is_fft_mode=True, dc_threshold_mv=threshold,
              highlight_masked=True, mask_color="magenta"))
    assert err == ""
    assert captured["bad_color"] == "magenta"
    # The export masks by value too (0 == the "no valid peak" sentinel, same
    # rule as the viewer's _redraw_image); on this fixture every
    # above-threshold pixel has a nonzero peak, so the value mask coincides
    # with the DC mask and the NaN set is exactly expect_masked.
    assert np.array_equal(np.isnan(captured["img"]), expect_masked)
    valid_vals = captured["img"][~expect_masked]
    assert not np.isnan(valid_vals).any() and (valid_vals != 0).all(), \
        "fixture precondition: every valid pixel has a nonzero peak"

    captured2 = _spy_draw(monkeypatch)
    err, _ = export_view_image(
        str(path), out_dir=str(out_dir), angle_idx=0, ch_idx=CH1_IDX,
        **_kw(is_fft_mode=True, dc_threshold_mv=threshold, highlight_masked=False))
    assert err == ""
    assert captured2["bad_color"] is None
    assert not np.isnan(captured2["img"]).any()


def test_auto_scale_uses_per_file_min_max(tmp_path, monkeypatch):
    path = tmp_path / "scale.sras"
    gen.write(path, n_angles=1, seed=8, samples_per_frame=64)
    out_dir = tmp_path / "out"
    out_dir.mkdir()

    captured = _spy_draw(monkeypatch)
    err, _ = export_view_image(
        str(path), out_dir=str(out_dir), angle_idx=0, ch_idx=CH4_IDX,
        **_kw(auto_scale=True))
    assert err == ""
    img = captured["img"]
    assert captured["vmin"] == pytest.approx(float(np.nanmin(img)))
    assert captured["vmax"] == pytest.approx(float(np.nanmax(img)))

    captured2 = _spy_draw(monkeypatch)
    err, _ = export_view_image(
        str(path), out_dir=str(out_dir), angle_idx=0, ch_idx=CH4_IDX,
        **_kw(auto_scale=False, vmin=-5.0, vmax=5.0))
    assert err == ""
    assert captured2["vmin"] == -5.0
    assert captured2["vmax"] == 5.0


# ---------------------------------------------------------------------------
# GUI dispatch: SrasViewerWindow._on_batch_export_images end-to-end
# ---------------------------------------------------------------------------

def _make_window(path) -> SrasViewerWindow:
    app = QApplication.instance() or QApplication([])   # noqa: F841
    win = SrasViewerWindow()
    win.show()
    win._load_file(str(path))
    assert wait_until(lambda: win._sras is not None), "file loaded"
    assert wait_until(lambda: all((a, CH4_IDX) in win._dc_cache
                                  for a in range(win._sras.n_angles))), \
        "DC precompute finished"
    return win


def test_batch_export_images_writes_one_png_per_file(tmp_path):
    path = tmp_path / "src.sras"
    gen.write(path, n_angles=2, seed=10, samples_per_frame=64)
    paths = [str(path)]
    for i in range(2):
        p2 = tmp_path / f"other{i}.sras"
        gen.write(p2, n_angles=2, seed=20 + i, samples_per_frame=64)
        paths.append(str(p2))
    out_dir = tmp_path / "images"
    out_dir.mkdir()

    win = _make_window(path)
    try:
        with patch("sras_viewer.main_window.QFileDialog.getOpenFileNames",
                   return_value=(paths, "")), \
             patch("sras_viewer.main_window.QFileDialog.getExistingDirectory",
                   return_value=str(out_dir)):
            win._on_batch_export_images()
        assert wait_until(lambda: not win._job_running("batch"), 60000), "batch ran"

        angle = win.spin_angle.value()
        ch_name = CH_NAMES[win.combo_channel.currentIndex()]
        expected_names = {f"{Path(p).stem}_angle{angle}_{ch_name}.png" for p in paths}
        actual_names = {p.name for p in out_dir.iterdir()}
        assert actual_names == expected_names
        assert "Batch export: 3/3 image(s)" in win.statusBar().currentMessage()
    finally:
        win.close()
        pump(200)


def test_batch_export_out_of_range_angle_reports_error_continues(tmp_path):
    good_path = tmp_path / "good.sras"
    short_path = tmp_path / "short.sras"
    gen.write(good_path, n_angles=3, seed=30, samples_per_frame=64)
    gen.write(short_path, n_angles=1, seed=31, samples_per_frame=64)
    out_dir = tmp_path / "images"
    out_dir.mkdir()

    win = _make_window(good_path)
    try:
        win.spin_angle.setValue(2)   # valid for good_path, out of range for short_path
        with patch("sras_viewer.main_window.QFileDialog.getOpenFileNames",
                   return_value=([str(good_path), str(short_path)], "")), \
             patch("sras_viewer.main_window.QFileDialog.getExistingDirectory",
                   return_value=str(out_dir)):
            win._on_batch_export_images()
        assert wait_until(lambda: not win._job_running("batch"), 60000), "batch ran"

        msg = win.statusBar().currentMessage()
        assert "Batch export: 1/2 image(s)" in msg, msg
        assert "1 failed" in msg, msg
        assert len(list(out_dir.iterdir())) == 1, \
            "the batch must not abort — the good file still exports"
    finally:
        win.close()
        pump(200)


def test_batch_export_busy_guard_skips_dialogs(tmp_path, monkeypatch):
    path = tmp_path / "busy.sras"
    gen.write(path, n_angles=1, seed=40, samples_per_frame=64)
    win = _make_window(path)
    try:
        monkeypatch.setattr(win, "_job_running", lambda key: True)
        with patch("sras_viewer.main_window.QFileDialog.getOpenFileNames") as mock_dlg:
            win._on_batch_export_images()
        assert mock_dlg.call_count == 0, \
            "the Jobs.BATCH busy guard must return before opening any dialog"
    finally:
        win.close()
        pump(200)


def test_batch_export_ignores_aligned_view_toggle(tmp_path):
    """Aligned View is geometry specific to whichever single file the
    Alignment Wizard last ran against and cannot be meaningfully applied
    across a batch of different files -- _on_batch_export_images must not
    read chk_aligned_view / self._alignment_result at all, regardless of
    what's checked in the live view."""
    path = tmp_path / "aligned.sras"
    gen.write(path, n_angles=1, seed=41, samples_per_frame=64)
    out_dir = tmp_path / "images"
    out_dir.mkdir()

    win = _make_window(path)
    try:
        captured_kwargs = []
        orig_init = BatchExportImagesWorker.__init__

        def spy_init(self, paths, **kw):
            captured_kwargs.append(kw)
            return orig_init(self, paths, **kw)

        win.chk_aligned_view.setChecked(True)
        with patch.object(BatchExportImagesWorker, "__init__", spy_init), \
             patch("sras_viewer.main_window.QFileDialog.getOpenFileNames",
                   return_value=([str(path)], "")), \
             patch("sras_viewer.main_window.QFileDialog.getExistingDirectory",
                   return_value=str(out_dir)):
            win._on_batch_export_images()
        assert wait_until(lambda: not win._job_running("batch"), 60000), "batch ran"

        assert len(captured_kwargs) == 1
        assert not any("align" in k.lower() for k in captured_kwargs[0]), \
            captured_kwargs[0].keys()
    finally:
        win.close()
        pump(200)


def test_batch_export_filename_collision_note(tmp_path):
    dir_a, dir_b = tmp_path / "dir_a", tmp_path / "dir_b"
    dir_a.mkdir()
    dir_b.mkdir()
    path_a, path_b = dir_a / "dup.sras", dir_b / "dup.sras"
    gen.write(path_a, n_angles=1, seed=50, samples_per_frame=64)
    gen.write(path_b, n_angles=1, seed=51, samples_per_frame=64)
    out_dir = tmp_path / "images"
    out_dir.mkdir()

    win = _make_window(path_a)
    try:
        with patch("sras_viewer.main_window.QFileDialog.getOpenFileNames",
                   return_value=([str(path_a), str(path_b)], "")), \
             patch("sras_viewer.main_window.QFileDialog.getExistingDirectory",
                   return_value=str(out_dir)):
            win._on_batch_export_images()
        assert wait_until(lambda: not win._job_running("batch"), 60000), "batch ran"

        msg = win.statusBar().currentMessage()
        assert "Batch export: 2/2 image(s)" in msg, msg
        assert "collision" in msg, msg
        assert len(list(out_dir.iterdir())) == 1, \
            "same-stem inputs silently overwrite to one output file"
    finally:
        win.close()
        pump(200)


def test_batch_export_does_not_modify_source_files(tmp_path):
    path = tmp_path / "untouched.sras"
    gen.write(path, n_angles=1, seed=60, samples_per_frame=64)
    before = path.read_bytes()
    out_dir = tmp_path / "images"
    out_dir.mkdir()

    win = _make_window(path)
    try:
        with patch("sras_viewer.main_window.QFileDialog.getOpenFileNames",
                   return_value=([str(path)], "")), \
             patch("sras_viewer.main_window.QFileDialog.getExistingDirectory",
                   return_value=str(out_dir)):
            win._on_batch_export_images()
        assert wait_until(lambda: not win._job_running("batch"), 60000), "batch ran"
    finally:
        win.close()
        pump(200)

    assert path.read_bytes() == before, "export must never write to the source file"