noaa-goes-visualization/tests/test_comparison_video.py

"""Tests for assembling the three-pane comparison video.

The property that matters is *alignment*: a variant that produced no composite for a
timestamp -- which is exactly what the current pipeline does past its gap limit --
must still occupy that slot, or its pane slides out of step with the others and the
comparison is meaningless.
"""

import os

import numpy as np
import pytest
from PIL import Image

import make_comparison_video as mcv
from suvi import paths

T0 = 1715299200


def write_stream(root, variant, satellite, timestamps, size=(64, 36)):
    directory = os.path.join(root, variant, f"goes{satellite}")
    os.makedirs(directory, exist_ok=True)
    for when in timestamps:
        img = Image.fromarray(np.full((size[1], size[0], 3), 40, dtype="uint8"))
        img.save(os.path.join(directory, f"Composite-{when}.jpg"), quality=90)
    return directory


def test_stream_frames_reads_timestamps(tmp_path):
    times = [T0, T0 + paths.CADENCE, T0 + 2 * paths.CADENCE]
    directory = write_stream(str(tmp_path), "pristine", 16, times)
    frames = mcv.stream_frames(directory)
    assert sorted(frames) == times
    assert all(os.path.exists(p) for p in frames.values())


def test_stream_frames_ignores_other_files(tmp_path):
    directory = write_stream(str(tmp_path), "new", 16, [T0])
    open(os.path.join(directory, "notes.txt"), "w").write("x")
    open(os.path.join(directory, "Composite-nonsense.jpg"), "w").write("x")
    assert list(mcv.stream_frames(directory)) == [T0]


def test_stream_frames_on_a_missing_directory(tmp_path):
    assert mcv.stream_frames(str(tmp_path / "absent")) == {}


def test_every_known_variant_is_labelled():
    """Variants are now a CLI choice; every name the tool accepts needs a label."""
    import render_real

    assert set(mcv.LABELS) == {"pristine", "today", "new", "learned", "damaged",
                               "optical_flow"}
    assert set(render_real.FILLERS) <= set(mcv.LABELS)     # its stream directories
    assert all(mcv.LABELS.values())


def test_encode_emits_one_frame_per_slot_including_gaps(tmp_path, monkeypatch):
    """A pane with holes must still be full length, or the panes drift apart.

    'today' legitimately has none for long gaps; those slots have to become black
    frames occupying the timeline, not vanish from it.
    """
    times = [T0, T0 + 2 * paths.CADENCE]  # slot 1 deliberately missing
    directory = write_stream(str(tmp_path), "today", 16, times)
    frames = mcv.stream_frames(directory)
    timeline = [T0 + i * paths.CADENCE for i in range(3)]

    written = []

    class FakeProcess:
        def __init__(self):
            self.stdin = self

        def write(self, data):
            written.append(("real", len(data)))

        def save(self, *a, **k):
            pass

        def close(self):
            pass

        def wait(self):
            return 0

    fake = FakeProcess()
    monkeypatch.setattr(mcv.subprocess, "Popen", lambda *a, **k: fake)

    class CountingImage:
        size = (64, 36)

        @staticmethod
        def fromarray(arr):
            class Blank:
                def save(self, handle, *a, **k):
                    written.append(("black", 0))
            return Blank()

        @staticmethod
        def open(path):
            return CountingImage

    monkeypatch.setattr(mcv, "Image", CountingImage)

    present = mcv.encode_stream("ffmpeg", frames, timeline, "out.mp4", 60, 18)
    assert present == 2
    assert len(written) == 3, "a missing slot did not occupy the timeline"
    assert [kind for kind, _ in written] == ["real", "black", "real"]


# ------------------------------------------------------ pristine reads the truth


def test_pristine_resolves_deleted_slots_from_the_archive(tmp_path):
    """The ground-truth pane must show the window as it really is.

    A slot the case deleted has no overlay path.  Resolving pristine through the
    overlay silently drops those timestamps -- 221 of them in one run -- leaving the
    reference pane shorter than the panes it exists to be compared against.
    """
    import bench
    from suvi import cases

    slot = (16, 171, T0)
    overlay = cases.Overlay(
        archive={slot: ("/archive/real.fits", None)},
        overrides={},
        deleted=frozenset({slot}),
    )
    assert overlay.path(slot) is None, "precondition: the case deleted this slot"
    assert overlay.truth_path(slot) == "/archive/real.fits"

    reads = []

    def fake_read(path):
        reads.append(path)
        return "pixels", None

    original = bench.fitsio.read_image
    bench.fitsio.read_image = fake_read
    try:
        cache = {}
        assert bench._resolve_band(overlay, slot, cache, use_truth=True) == "pixels"
        assert reads == ["/archive/real.fits"]
        # Without use_truth the slot is genuinely gone, which is right for the
        # variants that are meant to show the damage.
        assert bench._resolve_band(overlay, slot, cache, use_truth=False) is None
    finally:
        bench.fitsio.read_image = original


def test_resolve_band_caches_truth_and_overlay_separately(tmp_path):
    """One cache serving both would return the wrong frame for one of them."""
    import bench
    from suvi import cases

    slot = (16, 171, T0)
    overlay = cases.Overlay(
        archive={slot: ("/archive/real.fits", None)},
        overrides={slot: "/overlay/damaged.fits"},
    )
    original = bench.fitsio.read_image
    bench.fitsio.read_image = lambda path: (path, None)
    try:
        cache = {}
        assert bench._resolve_band(overlay, slot, cache, use_truth=True) == "/archive/real.fits"
        assert bench._resolve_band(overlay, slot, cache, use_truth=False) == "/overlay/damaged.fits"
    finally:
        bench.fitsio.read_image = original


def test_pane_spec_carries_per_pane_satellites():
    assert mcv.parse_panes("damaged:18,new,damaged,learned", 16, "/r") == (
        ("damaged", 18, "/r"), ("new", 16, "/r"), ("damaged", 16, "/r"),
        ("learned", 16, "/r"),
    )


def test_a_pane_can_come_from_another_render_root():
    """A new model renders only its learned stream; the inputs and the optical
    flow repair of the same case are read from the earlier render."""
    panes = mcv.parse_panes("damaged:18,new,damaged,learned@/renders/v7_s128",
                            16, "/renders/v7")
    assert panes[3] == ("learned", 16, "/renders/v7_s128")
    assert panes[0] == ("damaged", 18, "/renders/v7")


def test_a_pane_from_another_root_names_it_on_screen():
    """Two models' outputs must never be mistakable for one another."""
    own = mcv.pane_label("learned", 16, "/renders/v7_s128/", "/renders/v7")
    shared = mcv.pane_label("learned", 16, "/renders/v7", "/renders/v7")
    assert own.endswith("| v7_s128")
    assert "v7" not in shared.split("|")[-1]


def test_panes_from_different_roots_with_the_same_variant_stay_distinct():
    panes = mcv.parse_panes("learned@/a,learned@/b", 16, "/r")
    assert len(set(panes)) == 2


def test_pane_spec_rejects_unknown_variants():
    with pytest.raises(SystemExit, match="unknown variant"):
        mcv.parse_panes("pristine,upscaled", 16)


# ------------------------------------------------------------------ layout


def _run_graph(tmp_path, panes, labels=None):
    """Run the real filter graph through ffmpeg on tiny clips; return the output size."""
    import subprocess
    ffmpeg = mcv.find_ffmpeg()
    clips = []
    for index in range(len(panes)):
        clip = str(tmp_path / f"pane{index}.mp4")
        subprocess.run(f'{ffmpeg} -y -loglevel error -f lavfi -i testsrc=size=192x108:rate=4 '
                       f'-t 1 -pix_fmt yuv420p "{clip}"', shell=True, check=True)
        clips.append(clip)
    font = os.path.join(os.path.dirname(os.path.abspath(mcv.__file__)), "OpenSans-Regular.ttf")
    out = str(tmp_path / "out.mp4")
    inputs = " ".join(f'-i "{c}"' for c in clips)
    graph = mcv.filter_graph(panes, "/r", font, labels)
    result = subprocess.run(f'{ffmpeg} -y -loglevel error {inputs} -filter_complex "{graph}" '
                            f'-map "[out]" -pix_fmt yuv420p "{out}"',
                            shell=True, capture_output=True, text=True)
    assert result.returncode == 0, result.stderr
    probe = subprocess.run(f'{ffmpeg} -i "{out}"', shell=True, capture_output=True, text=True)
    import re
    return tuple(int(v) for v in re.search(r", (\d+)x(\d+)", probe.stderr).groups())


def test_a_single_stream_is_labelled_at_full_size_without_stacking(tmp_path):
    """The real-data run renders one stream; ffmpeg's stack filters refuse a
    single input, which is what failed the first GOES-16 year video."""
    graph = mcv.filter_graph((("learned", 16, "/r"),), "/r", "font.ttf")
    assert "stack" not in graph and "scale=" not in graph
    assert _run_graph(tmp_path, (("learned", 16, "/r"),)) == (192, 108)


def test_two_panes_sit_side_by_side_at_half_width(tmp_path):
    panes = (("new", 16, "/r"), ("learned", 16, "/r"))
    assert "hstack=inputs=2" in mcv.filter_graph(panes, "/r", "font.ttf")
    assert _run_graph(tmp_path, panes) == (1920, 540)


def test_four_panes_form_a_two_by_two_grid(tmp_path):
    panes = (("damaged", 18, "/r"), ("new", 16, "/r"), ("damaged", 16, "/r"),
             ("learned", 16, "/r"))
    assert "xstack=inputs=4:layout=0_0|w0_0|0_h0|w0_h0" in mcv.filter_graph(panes, "/r", "f")
    assert _run_graph(tmp_path, panes) == (1920, 1080)


# ------------------------------------------------------------------ labels


def test_without_labels_the_defaults_stand():
    assert mcv.parse_labels(None, 4) is None


def test_labels_replace_the_default_pane_labels():
    panes = (("optical_flow", 16, "/a"), ("learned", 16, "/b"))
    labels = mcv.parse_labels("OLD DETECTION + OPTICAL FLOW, NEW DETECTION + v8", 2)
    assert labels == ("OLD DETECTION + OPTICAL FLOW", "NEW DETECTION + v8")
    graph = mcv.filter_graph(panes, "/r", "font.ttf", labels)
    assert "text='OLD DETECTION + OPTICAL FLOW'" in graph
    assert "text='NEW DETECTION + v8'" in graph
    assert "LEARNED" not in graph


def test_labels_sit_bottom_centre_clear_of_the_composite_title():
    """At the top they hid the timestamp a frame is found again by."""
    graph = mcv.filter_graph((("learned", 16, "/r"), ("new", 16, "/r")), "/r", "f")
    assert graph.count("x=(w-text_w)/2:y=h-text_h-12") == 2
    assert "y=12" not in graph


def test_one_label_per_pane_is_required():
    with pytest.raises(SystemExit, match="names 2 panes; there are 4"):
        mcv.parse_labels("a,b", 4)


@pytest.mark.parametrize("label", [
    "it's", 'say "hi"', "a:b", "back\\slash", "$(rm -rf ~)", "`id`", "a;b", "a[0]",
    "", "   ", "new\nline",
])
def test_labels_that_would_break_the_filter_or_the_shell_are_refused(label):
    with pytest.raises(SystemExit, match="label"):
        mcv.parse_labels(f"ok,{label}", 2)


def test_every_allowed_label_character_survives_ffmpeg(tmp_path):
    panes = (("optical_flow", 16, "/r"), ("learned", 16, "/r"))
    labels = mcv.parse_labels("A-z 0.9 + (x) | a/b_c, OLD DETECTION + OPTICAL FLOW", 2)
    assert _run_graph(tmp_path, panes, labels) == (1920, 540)


# ----------------------------------------------------------------- periods


JAN1 = 1704067200                       # 2024-01-01 00:00 UTC


def test_periods_parse_as_utc_and_come_back_in_time_order():
    periods = mcv.parse_periods("2024-10-03..2024-10-06, 2024-01-01..2024-01-08")
    assert periods == ((JAN1, JAN1 + 7 * 86400), (1727913600, 1727913600 + 3 * 86400))


def test_adjacent_periods_are_allowed_but_overlapping_ones_are_not():
    assert len(mcv.parse_periods("2024-01-01..2024-01-02,2024-01-02..2024-01-03")) == 2
    with pytest.raises(SystemExit, match="overlap"):
        mcv.parse_periods("2024-01-01..2024-01-05,2024-01-04..2024-01-08")


@pytest.mark.parametrize("spec", ["2024-01-01", "2024-01-01..", "..2024-01-02",
                                  "2024-13-01..2024-12-01", "a..b",
                                  "2024-01-01..2024-01-02..2024-01-03", ""])
def test_malformed_periods_are_refused(spec):
    with pytest.raises(SystemExit, match="expected YYYY-MM-DD"):
        mcv.parse_periods(spec)


def test_a_period_that_ends_before_it_starts_is_refused():
    with pytest.raises(SystemExit, match="ends before"):
        mcv.parse_periods("2024-01-02..2024-01-01")
    with pytest.raises(SystemExit, match="ends before"):
        mcv.parse_periods("2024-01-02..2024-01-02")


def test_the_default_timeline_spans_first_to_last_composite_with_gaps_kept():
    step = paths.CADENCE
    streams = {"a": {JAN1: "x", JAN1 + 3 * step: "y"}, "b": {JAN1 + 5 * step: "z"}}
    assert mcv.build_timeline(streams) == [JAN1 + i * step for i in range(6)]
    assert mcv.build_timeline({"a": {}}) == []


def test_periods_keep_only_their_own_slots_back_to_back():
    """A whole-year pane must not stretch three sample weeks into a year."""
    step = paths.CADENCE
    streams = {"year": {JAN1 + i * step: "x" for i in range(0, 360 * 366, 50)}}
    periods = ((JAN1, JAN1 + 3 * step), (JAN1 + 86400, JAN1 + 86400 + 2 * step))
    assert mcv.build_timeline(streams, periods) == [
        JAN1, JAN1 + step, JAN1 + 2 * step, JAN1 + 86400, JAN1 + 86400 + step]