orisys-app/shared/tests/test_depth_replay.py

"""Offline depth replay: where to start, how far to step, which chain to run.

P1-3 and P1-4: the pane kept showing the previous frame's depth after the cursor
moved, the recomputer re-read ``target`` frames on every call (so requesting
frames 2 and 3 landed on 5), it skipped the live processing chain, and it
reduced the recorded config path to a name that could load a different file.
"""

from __future__ import annotations

import shutil
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import patch

import numpy as np
import pytest

from orisys_app_shared.record.depth import (
    STATE_APPROXIMATE,
    STATE_EXACT,
    STATE_PREROLL,
    DepthRecomputer,
)

SCRATCH = Path(__file__).resolve().parent / "_scratch"


def _root() -> Path:
    shutil.rmtree(SCRATCH, ignore_errors=True)
    SCRATCH.mkdir(parents=True, exist_ok=True)
    return SCRATCH


class _CountingSensor:
    """A sensor whose reported depth is the number of frames it has read."""

    opened: list[dict] = []

    def __init__(self, *args, **kwargs):
        _CountingSensor.opened.append({"args": args, "kwargs": kwargs})
        self.n = 0
        self.chain_calls = 0
        self.initial_img = "reference-frame"
        # The real Sensor opens its source in the constructor, so the first
        # source is part of its history too.
        self.vid_src = str(args[0]) if args else ""
        self.sources: list[str] = [self.vid_src] if self.vid_src else []
        self.resets = 0
        self.info = SimpleNamespace(DEPTH=1)

    def get_img(self):
        self.n += 1
        return np.ones((2, 2))

    def compute_deformation(self, **_kwargs):
        self.chain_calls += 1

    def compute_contact(self):
        self.chain_calls += 1

    def reset(self):
        self.resets += 1

    def read_info(self, _channel):
        return np.array([self.n])

    def _open_video_source(self, _override, _box):
        self.sources.append(str(self.vid_src))
        # A real reopen restarts the file from its first frame.
        self.n = 0

    def disconnect(self):
        return None


def _reader(session: dict, *, segments=("raw_0001.mp4",), frames=()):
    directory = _root()
    for name in segments:
        (directory / name).write_bytes(b"")
    return SimpleNamespace(
        session=lambda: session,
        segment_path=lambda segment: directory / f"raw_{segment:04d}.mp4",
        segments=lambda: [directory / name for name in segments],
        info=lambda: SimpleNamespace(frames=list(frames)),
        frame_of=lambda frame_id: next(
            (frame for frame in frames if frame.frame_id == frame_id), None
        ),
        frame_at=lambda t: None,
        last_reset_frame=lambda frame_id: None,
        segment_start_frame=lambda segment: next(
            (frame for frame in frames if frame.segment == segment), None
        ),
        preroll_frames=lambda: 0,
        preroll_segment=lambda: None,
        directory=directory,
    )


def _session(config_ref: str) -> dict:
    return {"device": {"config_ref": config_ref, "sensor_type": "box"}}


def _config_file() -> str:
    """A config that exists, so ``preflight`` is not what a test measures."""
    path = _root() / "sensor.json"
    path.write_text("{}", encoding="utf-8")
    return str(path)


def _reader_with_config(frames):
    """A reader whose recorded config really exists on disk.

    ``_reader`` recreates the scratch directory, so the config file is written
    *after* the reader is built.
    """
    directory = _root()
    config = directory / "sensor.json"
    reader = _reader(_session(str(config)), frames=frames)
    config.write_text("{}", encoding="utf-8")
    return reader


def test_stepping_continues_from_the_current_position():
    """The old code re-read ``target`` frames every call: 2 then 3 landed on 5."""
    _CountingSensor.opened.clear()
    computer = DepthRecomputer(_reader(_session("/cfg/sensor.json")))
    with patch("orisys.Sensor", _CountingSensor):
        first = computer._depth_from_segment(Path("raw_0001.mp4"), 2)
        second = computer._depth_from_segment(Path("raw_0001.mp4"), 3)
    assert first.tolist() == [2]
    assert second.tolist() == [3]
    # One sensor for the whole replay: reopening would reset the algorithm state.
    assert len(_CountingSensor.opened) == 1
    computer.close()


def test_stepping_backwards_reopens_the_source_but_keeps_the_processor():
    _CountingSensor.opened.clear()
    computer = DepthRecomputer(_reader(_session("/cfg/sensor.json")))
    with patch("orisys.Sensor", _CountingSensor):
        assert computer._depth_from_segment(Path("raw_0001.mp4"), 5).tolist() == [5]
        # Reading an earlier frame restarts the file; the processor object (and
        # therefore the tracker) is not thrown away.
        assert computer._depth_from_segment(Path("raw_0001.mp4"), 2).tolist() == [2]
        sensor = computer._sensor._sensor
    assert len(_CountingSensor.opened) == 1, "the processor was recreated"
    assert len(sensor.sources) == 2
    assert all(source.endswith("raw_0001.mp4") for source in sensor.sources)
    computer.close()


def test_the_live_processing_chain_runs_for_every_replayed_frame():
    """Without it the replay is a different algorithm than the recording."""
    _CountingSensor.opened.clear()
    computer = DepthRecomputer(_reader(_session("/cfg/sensor.json")))
    with patch("orisys.Sensor", _CountingSensor):
        computer._depth_from_segment(Path("raw_0001.mp4"), 3)
        sensor = computer._sensor._sensor
    assert sensor.chain_calls == 6, "one deformation + one contact per frame"
    computer.close()


def test_the_exact_config_path_is_opened(tmp_path=None):
    """A reduced name could load a different file with the same name."""
    _CountingSensor.opened.clear()
    config = Path("D:/custom/box/sensor.json")
    computer = DepthRecomputer(_reader(_session(str(config))))
    with patch("orisys.Sensor", _CountingSensor):
        computer._depth_from_segment(Path("raw_0001.mp4"), 1)
    assert _CountingSensor.opened[0]["kwargs"]["config_name"] == str(config)
    computer.close()


def test_the_plan_starts_at_the_last_reset_in_the_same_segment():
    frames = [
        SimpleNamespace(
            frame_id=index,
            t_rel_s=index / 30.0,
            segment=1,
            segment_frame=index + 1,
            in_video=True,
            dropped=False,
        )
        for index in range(10)
    ]
    reader = _reader(_session("/cfg/sensor.json"), frames=frames)
    reader.last_reset_frame = lambda frame_id: 6
    computer = DepthRecomputer(reader)
    plan = computer.plan_for(frames[8])
    assert plan.stop_at == 9
    assert plan.anchor == "reset"
    assert plan.state == STATE_EXACT
    assert plan.reset_frame_id == 6
    assert len(plan.steps) == 1
    step = plan.steps[0]
    assert step.source == reader.directory / "raw_0001.mp4"
    assert step.to_frame == 9
    assert step.reset_after == 7, "the reset boundary is not re-applied by the plan"
    assert plan.frames_to_read == 9
    assert plan.preroll is None
    assert plan.warnings == ()
    computer.close()


def test_a_reset_that_is_not_in_the_video_is_not_a_recovered_state():
    """Finding the event is not the same as being able to start there."""
    frames = [
        SimpleNamespace(
            frame_id=index,
            t_rel_s=index / 30.0,
            segment=1,
            segment_frame=index + 1,
            in_video=index != 6,  # the reset frame never reached the video
            dropped=False,
        )
        for index in range(10)
    ]
    reader = _reader(_session("/cfg/sensor.json"), frames=frames)
    reader.last_reset_frame = lambda frame_id: 6
    computer = DepthRecomputer(reader)
    plan = computer.plan_for(frames[8])
    assert plan.anchor != "reset"
    assert plan.state == STATE_APPROXIMATE
    assert "reset_frame_not_in_video" in plan.warnings
    assert "state_approximate" in plan.warnings
    computer.close()


def test_the_plan_uses_the_pre_roll_for_the_start_of_a_recording():
    frames = [
        SimpleNamespace(
            frame_id=index,
            t_rel_s=index / 30.0,
            segment=1,
            segment_frame=index + 1,
            in_video=True,
            dropped=False,
        )
        for index in range(5)
    ]

    def reader_factory():
        reader = _reader(_session("/cfg/sensor.json"), frames=frames)
        reader.preroll_frames = lambda: 30
        reader.preroll_segment = lambda: reader.directory / "preroll_0001.mp4"
        return reader

    reader = reader_factory()
    (reader.directory / "preroll_0001.mp4").write_bytes(b"")
    computer = DepthRecomputer(reader)
    plan = computer.plan_for(frames[3])
    assert plan.preroll is not None
    assert plan.preroll_frames == 30
    assert plan.anchor == "preroll"
    assert plan.state == STATE_PREROLL
    assert plan.warnings == ()
    computer.close()


def test_a_segment_without_history_is_flagged_approximate():
    """Starting mid-recording with no reset and no pre-roll cannot be exact."""
    frames = [
        SimpleNamespace(
            frame_id=index,
            t_rel_s=index / 30.0,
            segment=2,
            segment_frame=index + 1,
            in_video=True,
            dropped=False,
        )
        for index in range(5)
    ]
    reader = _reader(_session("/cfg/sensor.json"), frames=frames)
    computer = DepthRecomputer(reader)
    plan = computer.plan_for(frames[2])
    assert plan.state == STATE_APPROXIMATE
    assert plan.anchor == "segment_start"
    assert "state_approximate" in plan.warnings
    computer.close()


def test_a_reset_in_an_earlier_segment_is_carried_across_by_swapping_the_source():
    """The processor survives the swap, so the state *can* cross a segment."""
    frames = [
        SimpleNamespace(
            frame_id=index,
            t_rel_s=index / 30.0,
            segment=1 if index < 3 else 2,
            segment_frame=index + 1 if index < 3 else index - 2,
            in_video=True,
            dropped=False,
        )
        for index in range(6)
    ]
    reader = _reader(
        _session("/cfg/sensor.json"),
        segments=("raw_0001.mp4", "raw_0002.mp4"),
        frames=frames,
    )
    reader.last_reset_frame = lambda frame_id: 1
    computer = DepthRecomputer(reader)
    plan = computer.plan_for(frames[5])
    assert plan.anchor == "reset"
    assert plan.state == STATE_EXACT
    assert "state_across_segment" not in plan.warnings
    assert [step.source.name for step in plan.steps] == ["raw_0001.mp4", "raw_0002.mp4"]
    assert plan.steps[0].reset_after == 2
    assert plan.steps[1].to_frame == 3
    computer.close()


def test_a_replay_that_would_read_too_far_back_is_reported_not_attempted():
    frames = [
        SimpleNamespace(
            frame_id=index,
            t_rel_s=index / 30.0,
            segment=1,
            segment_frame=index + 1,
            in_video=True,
            dropped=False,
        )
        for index in range(20)
    ]
    reader = _reader(_session("/cfg/sensor.json"), frames=frames)
    reader.last_reset_frame = lambda frame_id: 1
    computer = DepthRecomputer(reader)
    with patch("orisys_app_shared.record.depth.MAX_REPLAY_FRAMES", 5):
        plan = computer.plan_for(frames[10])
    assert "state_too_far_back" in plan.warnings
    assert plan.state == STATE_APPROXIMATE
    computer.close()


def test_the_same_frame_is_answered_from_the_cache():
    """A second request for the frame just read is a cache hit, not None."""
    _CountingSensor.opened.clear()
    computer = DepthRecomputer(_reader(_session("/cfg/sensor.json")))
    with patch("orisys.Sensor", _CountingSensor):
        first = computer._depth_from_segment(Path("raw_0001.mp4"), 3)
        second = computer._depth_from_segment(Path("raw_0001.mp4"), 3)
        sensor = computer._sensor._sensor
    assert first.tolist() == [3]
    assert second is not None and second.tolist() == [3]
    assert sensor.n == 3, "the cached answer re-read the stream"
    computer.close()


def test_the_processor_survives_a_pre_roll_to_segment_switch():
    """Only the video source is replaced; the tracker state must survive."""
    _CountingSensor.opened.clear()
    reader = _reader(_session("/cfg/sensor.json"))
    preroll = reader.directory / "preroll_0001.mp4"
    preroll.write_bytes(b"")
    computer = DepthRecomputer(reader)
    with patch("orisys.Sensor", _CountingSensor):
        computer._feed(preroll, 2)
        sensor = computer._sensor._sensor
        assert sensor.n == 2, "the pre-roll frames were not fed"
        # Feed the segment: same processor, new source, position restarted.
        computer._feed(reader.directory / "raw_0001.mp4", 1)
    assert len(_CountingSensor.opened) == 1, "the processor was recreated"
    assert computer._sensor._sensor is sensor
    assert len(sensor.sources) == 2
    assert sensor.sources[0].endswith("preroll_0001.mp4")
    assert sensor.sources[1].endswith("raw_0001.mp4")
    assert sensor.initial_img == "reference-frame", "the tracker reference was reset"
    computer.close()


def test_a_reset_boundary_is_re_applied_during_the_replay():
    """The live zeroing cleared the tracker; the replay must do the same."""
    _CountingSensor.opened.clear()
    frames = [
        SimpleNamespace(
            frame_id=index,
            t_rel_s=index / 30.0,
            segment=1,
            segment_frame=index + 1,
            in_video=True,
            dropped=False,
        )
        for index in range(10)
    ]
    reader = _reader_with_config(frames)
    reader.last_reset_frame = lambda frame_id: 6
    reader.frame_at = lambda _t: frames[8]
    computer = DepthRecomputer(reader)
    with patch("orisys.Sensor", _CountingSensor):
        result = computer.depth_at(frames[8].t_rel_s)
        sensor = computer._sensor._sensor
    assert result.ok, result.error
    assert sensor.resets == 1, "the reset boundary was not re-applied"
    assert result.warnings == (), result.warnings
    computer.close()


def test_a_dropped_frame_is_reported_not_substituted():
    """The index marks the hole; replay must say so instead of showing a neighbour."""
    frames = [
        SimpleNamespace(
            frame_id=0,
            t_rel_s=0.0,
            segment=1,
            segment_frame=1,
            in_video=True,
            dropped=False,
        ),
        SimpleNamespace(
            frame_id=1,
            t_rel_s=1 / 30.0,
            segment=1,
            segment_frame=0,
            in_video=False,
            dropped=True,
        ),
    ]
    reader = _reader_with_config(frames)
    reader.frame_at = lambda _t: frames[1]
    computer = DepthRecomputer(reader)
    result = computer.depth_at(1 / 30.0)
    assert not result.ok
    assert result.error == "frame_not_recorded"
    computer.close()