"""Acceptance for the single 2D image backend.

The point of this module is that live, recomputed, replayed and previewed frames go
through *one* renderer and that the view modes come from what the frame carries. These
tests pin both halves: a plan derived from fields (not from a view preset), and
pixel-for-pixel agreement with the primitives the application used before.
"""

from __future__ import annotations

import numpy as np
import pytest
from orisys_app_shared.acquisition import image_backend as backend
from orisys_app_shared.acquisition.arrows import (
    depth_map_to_bgr,
    draw_arrow_display,
    flow_colormap,
)


def _image(h: int = 12, w: int = 16) -> np.ndarray:
    return np.arange(h * w * 3, dtype=np.uint8).reshape(h, w, 3)


def _flow(h: int = 12, w: int = 16) -> np.ndarray:
    flow = np.zeros((h, w, 2), dtype=np.float32)
    flow[..., 0] = 3.0
    flow[..., 1] = -1.5
    return flow


# ------------------------------------------------------------------------- plans
def test_a_live_frame_offers_the_derived_modes_and_needs_no_note():
    plan = backend.plan_image(
        backend.ImageFrame(_image(), flow=_flow(), depth=np.ones((12, 16)))
    )
    assert plan.available_modes == ("arrows_pure", "arrows_stitched", "depth")
    assert plan.default_mode == "arrows_pure"
    assert plan.missing == ()
    assert plan.note_key == ""


def test_a_recorded_frame_is_raw_and_says_so():
    """No flow and no depth: the frame *is* the picture, and the note explains it."""
    plan = backend.plan_image(backend.ImageFrame(_image(), source=backend.SOURCE_REPLAY))
    assert plan.available_modes == ("recorded",)
    assert plan.default_mode == "recorded"
    assert plan.missing == ("flow", "depth")
    assert plan.note_key == backend.RAW_NOTE_KEY
    assert plan.source == "replay"


def test_a_derived_source_keeps_every_mode_and_reports_what_is_missing():
    """A dropped frame must not silently switch the operator's view.

    Live and recomputed frames may be missing a field for one frame; the mode stays
    selectable and ``missing`` says which field is absent, so the panel can state it
    instead of the view looking broken.
    """
    flow_only = backend.plan_image(backend.ImageFrame(_image(), flow=_flow()))
    assert flow_only.available_modes == backend.DERIVED_MODES
    assert flow_only.missing == ("depth",)
    assert flow_only.default_mode == "arrows_pure"
    assert flow_only.note_key == backend.MISSING_NOTE_KEY
    assert flow_only.note_fields == ("depth",)
    assert flow_only.needs("depth") == "depth"
    assert flow_only.needs("arrows_pure") is None

    depth_only = backend.plan_image(backend.ImageFrame(_image(), depth=np.ones((12, 16))))
    assert depth_only.available_modes == backend.DERIVED_MODES
    assert depth_only.missing == ("flow",)
    assert depth_only.default_mode == "depth"
    assert depth_only.needs("arrows_pure") == "flow"

    complete = backend.plan_image(
        backend.ImageFrame(_image(), flow=_flow(), depth=np.ones((12, 16)))
    )
    assert complete.missing == ()
    assert complete.note_key == ""


def test_a_raw_source_cannot_offer_a_derived_view_at_all():
    """A recording and a camera preview have no derived field, ever."""
    for source in (backend.SOURCE_REPLAY, backend.SOURCE_PREVIEW):
        plan = backend.plan_image(backend.ImageFrame(_image(), source=source))
        assert plan.available_modes == ("recorded",)
        assert plan.note_key == backend.RAW_NOTE_KEY
        assert plan.note_fields == ()


def test_the_plan_carries_the_frame_identity_for_the_panel():
    plan = backend.plan_image(
        backend.ImageFrame(
            _image(), source=backend.SOURCE_RECOMPUTE, frame_id=42, t_rel_s=1.25, stale=True
        )
    )
    assert (plan.source, plan.frame_id, plan.t_rel_s, plan.stale) == ("recompute", 42, 1.25, True)


def test_a_mode_the_plan_does_not_offer_falls_back_to_the_default():
    """Never draw a view the source cannot have at all."""
    recorded = backend.plan_image(
        backend.ImageFrame(_image(), source=backend.SOURCE_REPLAY)
    )
    assert recorded.resolve("arrows_pure") == "recorded"
    assert recorded.resolve("nonsense") == "recorded"
    live = backend.plan_image(backend.ImageFrame(_image(), flow=_flow()))
    assert live.resolve("recorded") == "arrows_pure", "a derived source has no raw mode"
    assert live.resolve("arrows_stitched") == "arrows_stitched"


# --------------------------------------------------------------- one renderer
def test_the_2d_canvas_draws_from_the_policy_it_is_given():
    """The threshold is the caller's — the pane's — not a per-view default.

    ``arrow_policy`` is the policy the 3D surface samples its overlay with
    (``visual_policy.policy_for_settings`` over the display settings). A sample inside
    the hysteresis band is the case that used to differ between the two figures, so it
    is the case pinned here: with a 0.4 px threshold nothing is drawn, with 0.3 px the
    same frame carries arrows.
    """
    from orisys_app_shared.scene3d.contracts import DisplaySettings
    from orisys_app_shared.scene3d.visual_policy import policy_for_settings

    h, w = 12, 16
    flow = np.zeros((h, w, 2), dtype=np.float32)
    flow[..., 0] = 0.175  # |2*flow| = 0.35 px: above flow_off_px, below flow_on_px
    frame = backend.ImageFrame(_image(h, w), flow=flow, flow_valid=True)
    base = np.zeros((h, w, 3), dtype=np.uint8)

    strict = policy_for_settings(DisplaySettings(flow_on_px=0.4, flow_off_px=0.3))
    loose = policy_for_settings(DisplaySettings(flow_on_px=0.3, flow_off_px=0.2))
    strict_image = backend.render_image(
        frame, mode="arrows_pure", draw_style="finger", arrow_policy=strict
    )
    loose_image = backend.render_image(
        frame, mode="arrows_pure", draw_style="finger", arrow_policy=loose
    )
    assert len(np.unique(strict_image.reshape(-1, 3), axis=0)) == 1, (
        "a 0.4 px threshold drew something for a 0.35 px sample"
    )
    assert len(np.unique(loose_image.reshape(-1, 3), axis=0)) > 1, (
        "a 0.3 px threshold drew nothing for a 0.35 px sample"
    )
    # The same policy governs the primitive directly, which is what the window calls
    # through the backend.
    direct = draw_arrow_display(base, flow, draw_style="finger", policy=strict)
    assert len(np.unique(direct.reshape(-1, 3), axis=0)) == 1


def test_the_2d_canvas_colours_one_magnitude_the_same_in_every_frame():
    """The colour range is fixed, so two frames are comparable by colour.

    The canvas used to normalise each frame against its own maximum: the same magnitude
    was one colour in a frame with a fast sample and another in a frame without one, so
    two frames could look alike while showing different motion. ``magnitude_range``
    defaults to :data:`arrows.FLOW_MAGNITUDE_RANGE`, and the old re-normalisation is
    available only when a caller asks for it.
    """
    from orisys_app_shared.acquisition.arrows import magnitude_to_ramp_index

    # One probe magnitude in two frames whose maxima differ by a factor of four.
    weak = magnitude_to_ramp_index(np.array([0.5, 1.0]))
    strong = magnitude_to_ramp_index(np.array([0.5, 4.0]))
    assert weak[0] == strong[0], (
        "the same 0.5 px magnitude took two different ramp positions because the frame's "
        f"maximum changed: {weak[0]} against {strong[0]}"
    )
    # The range is the documented default, not one frame's own span.
    assert magnitude_to_ramp_index(np.array([0.0, 2.0])).tolist() == [0, 255]
    # The per-frame mode is still available, and is what made the colours drift.
    assert magnitude_to_ramp_index(np.array([0.5, 4.0]), per_frame=True)[0] != weak[0]
    # A caller may pin the range it wants instead: 1.0 of a 0..4.0 range is a quarter up.
    assert magnitude_to_ramp_index(np.array([1.0]), value_range=(0.0, 4.0))[0] == 63


def test_the_2d_canvas_draws_the_same_magnitude_in_the_same_colour():
    """The painted canvas obeys the fixed range through :func:`draw_arrow_display`.

    The helper's contract is one thing; this is the pixel an operator looks at. Two
    frames carry the same magnitude at the same sample and different maxima, and the
    sample's pixel is the same colour in both.
    """
    from orisys_app_shared.acquisition.arrows import ARROW_BG_BGR, flow_colormap
    from orisys_app_shared.scene3d.visual_policy import VisualPolicy

    h, w = 12, 12
    policy = VisualPolicy(flow_on_px=0.4, flow_grid_spacing=1, flow_edge_margin=0)

    def canvas(probe: float, strongest: float) -> np.ndarray:
        flow = np.zeros((h, w, 2), dtype=np.float64)
        flow[8, 8] = (probe, 0.0)
        flow[2, 8] = (strongest, 0.0)
        return draw_arrow_display(
            np.zeros((h, w, 3), dtype=np.uint8),
            flow,
            draw_style="finger",
            policy=policy,
            colormap=flow_colormap("dark"),
        )

    def sample_colour(drawn: np.ndarray, row: int, col: int) -> tuple[int, ...]:
        """The colour this sample was drawn in.

        The canvas is one flat colour and this sample is drawn at the centre of the
        patch, so any pixel that is not the canvas is the drawn sample.
        """
        margin = 2
        patch = drawn[row - margin : row + margin + 1, col - margin : col + margin + 1]
        background = np.array(ARROW_BG_BGR, dtype=np.uint8)
        drawn_pixels = patch[np.any(patch != background, axis=-1)]
        assert drawn_pixels.size, f"nothing was drawn at {(row, col)}"
        return tuple(int(v) for v in drawn_pixels[0])

    quiet = canvas(0.25, 0.5)  # probe |2*flow| = 0.5 px, frame maximum 1.0 px
    loud = canvas(0.25, 4.0)  # same probe, frame maximum 8.0 px
    assert sample_colour(quiet, 8, 8) == sample_colour(loud, 8, 8), (
        "one magnitude took two colours because another sample moved faster"
    )


def test_a_failed_draw_is_not_an_idle_frame():
    """The base image comes back with a code on it, because the pixels cannot say so.

    A render failure and a frame with nothing to draw produced the same picture, so
    "the arrows are missing" was indistinguishable from "the device measured nothing".
    The code travels on the returned image; the caller surfaces it (this backend's own
    panel wiring is not part of this change).
    """
    from orisys_app_shared.acquisition.arrows import ARROW_RENDER_FAILED

    base = np.zeros((12, 12, 3), dtype=np.uint8)
    base[:, :] = (0x20, 0x30, 0x40)
    # ``select_flow_samples`` rejects a field that is not [H,W,2], which is a draw
    # failure rather than an empty measurement.
    failed = draw_arrow_display(base, np.zeros((12, 12), dtype=np.float32), draw_style="finger")
    assert np.array_equal(failed, base), "a failed draw must still hand back the image"
    assert failed.frame_diagnostic is not None
    assert failed.frame_diagnostic.code == ARROW_RENDER_FAILED
    assert failed.frame_diagnostic.severity == "warning"
    # The exception's own message is not carried: diagnostics travel to the UI and logs.
    assert "Traceback" not in failed.frame_diagnostic.message

    # A frame with no flow at all is a legitimate empty frame, not a failure.
    idle = draw_arrow_display(base, None, draw_style="finger")
    assert idle.frame_diagnostic is None
    assert np.array_equal(idle, base)

    # A drawn frame says nothing, and stays an ordinary array for every other caller.
    drawn = draw_arrow_display(
        base,
        np.full((12, 12, 2), 0.5, dtype=np.float32),
        draw_style="finger",
    )
    assert drawn.frame_diagnostic is None
    assert isinstance(drawn, np.ndarray)
    assert drawn.shape == base.shape and drawn.dtype == np.uint8


def test_the_backend_renders_what_the_primitives_rendered_before():
    """Live went through these two calls; the backend must not change a pixel."""
    image = _image()
    flow = _flow()
    mask = np.ones((12, 16), dtype=bool)
    mask[0, 0] = False

    frame = backend.ImageFrame(image, flow=flow, valid_mask=mask)
    plan = backend.plan_image(frame)
    for mode in ("arrows_pure", "arrows_stitched"):
        expected = draw_arrow_display(
            image,
            flow,
            background="captured" if mode == "arrows_stitched" else "pure_color",
            draw_style="finger",
            valid_mask=mask,
            colormap=flow_colormap("dark"),
        )
        got = backend.render_image(frame, plan, mode=mode, draw_style="finger")
        assert np.array_equal(got, expected), mode

    depth = np.linspace(0.0, 1.0, 12 * 16, dtype=np.float32).reshape(12, 16)
    depth_frame = backend.ImageFrame(image, depth=depth, valid_mask=mask)
    # The depth picture uses the *theme's* ramp — the same table the panel's colour bar
    # paints. The window used to leave this argument out, so the map was drawn with
    # OpenCV's JET while the legend under it showed the brand ramp: two colour tables
    # for one picture, and the legend claiming they were the same.
    expected_depth = depth_map_to_bgr(
        depth, flow_colormap("dark"), value_range=(0.0, 1.0), invalid_mask=mask
    )
    got_depth = backend.render_image(
        depth_frame, mode="depth", depth_range=(0.0, 1.0)
    )
    assert np.array_equal(got_depth, expected_depth)


def test_a_recorded_frame_is_passed_through_and_never_resized_to_fit():
    image = _image(9, 7)
    got = backend.render_image(
        backend.ImageFrame(image, source=backend.SOURCE_REPLAY), mode="recorded"
    )
    assert got is not None and got.shape == image.shape
    assert np.array_equal(got, image)


def test_nothing_to_draw_is_none_not_a_blank_picture():
    assert backend.render_image(backend.ImageFrame(), mode="recorded") is None
    assert backend.render_image(
        backend.ImageFrame(None, flow=_flow()), mode="arrows_pure"
    ) is None
    assert backend.render_image(backend.ImageFrame(None, depth=None), mode="depth") is None


def test_a_grey_base_frame_is_widened_once_by_the_backend():
    """A 2-D frame must reach the shader as BGR, and only one place may do it."""
    grey = np.zeros((8, 6), dtype=np.uint8)
    got = backend.render_image(backend.ImageFrame(grey), mode="recorded")
    assert got is not None and got.shape == (8, 6, 3)
    assert np.array_equal(got, np.stack([grey] * 3, axis=-1))


@pytest.mark.parametrize("mode", list(backend.ALL_MODES))
def test_every_offered_mode_renders_something_or_says_it_cannot(mode):
    """A mode the plan offers must never come back empty."""
    frame = backend.ImageFrame(_image(), flow=_flow(), depth=np.ones((12, 16)))
    plan = backend.plan_image(frame)
    if not plan.accepts(mode):
        assert plan.resolve(mode) in plan.available_modes
        return
    assert backend.render_image(frame, plan, mode=mode) is not None


def test_a_missing_field_is_reported_not_invented():
    """Depth with no depth says so; arrows with no flow show the base picture."""
    no_depth = backend.ImageFrame(_image(), flow=_flow())
    assert backend.render_image(no_depth, mode="depth") is None
    no_flow = backend.ImageFrame(_image(), depth=np.ones((12, 16)))
    base = backend.render_image(no_flow, mode="arrows_pure")
    assert base is not None
    assert np.array_equal(base, _image()), "the base picture must not be redrawn"


# ------------------------------------------------------- the latched boundary
#: A domain the same size and the same kind of shape as the shipped finger sensor's
#: 300x280 UV island: one large blob with a long smooth rim, so "a one-pixel wobble"
#: and "the island moved" have the overlap the real package has (measured on the
#: shipped island: a 1 px wobble scores IoU 0.9633, a 6 px move 0.9102).
DOMAIN_HW = (300, 280)
ISLAND_RADIUS_PX = 95


def _island(h: int = DOMAIN_HW[0], w: int = DOMAIN_HW[1]) -> np.ndarray:
    yy, xx = np.mgrid[0:h, 0:w]
    return ((yy - h / 2.0) ** 2 + (xx - w / 2.0) ** 2) <= float(ISLAND_RADIUS_PX) ** 2


def _blank_image(h: int = DOMAIN_HW[0], w: int = DOMAIN_HW[1]) -> np.ndarray:
    """A base picture the two views do not read: only depth and the mask are drawn."""
    return np.zeros((h, w, 3), dtype=np.uint8)


def _shift(mask: np.ndarray, dy: int, dx: int) -> np.ndarray:
    out = np.zeros_like(mask)
    h, w = mask.shape
    ys = slice(max(0, dy), h + min(0, dy))
    xs = slice(max(0, dx), w + min(0, dx))
    ys2 = slice(max(0, -dy), h + min(0, -dy))
    xs2 = slice(max(0, -dx), w + min(0, -dx))
    out[ys, xs] = mask[ys2, xs2]
    return out


def test_overlap_of_two_identical_masks_is_one():
    mask = _island()
    assert backend.mask_overlap(mask, mask) == 1.0
    assert backend.mask_overlap(mask, ~mask) == 0.0
    # Two frames that agree there is no island at all are not a change.
    empty = np.zeros((4, 4), dtype=bool)
    assert backend.mask_overlap(empty, empty) == 1.0


def test_the_gate_separates_a_rebuilt_mask_from_a_moved_one():
    """The threshold itself, on the two cases it exists to tell apart."""
    island = _island()
    for dy, dx in ((1, 0), (0, 1), (1, 1), (-1, 1), (2, 0)):
        assert backend.mask_overlap(island, _shift(island, dy, dx)) >= (
            backend.DOMAIN_MASK_MIN_IOU
        ), f"a {dy},{dx} rebuild would be taken for a moved boundary"
    assert backend.mask_overlap(island, _shift(island, 6, 0)) < backend.DOMAIN_MASK_MIN_IOU


def test_a_stable_domain_mask_is_latched_once_for_the_session():
    """The device's island is geometry: rebuilt per frame, decided once."""
    latch = backend.DomainMaskLatch()
    mask = _island()
    seen = {latch.update(mask).tobytes() for _ in range(20)}
    assert len(seen) == 1, "a stable mask was rebuilt into a different array"
    assert np.array_equal(latch.mask, mask)
    assert (latch.adoptions, latch.held) == (1, 19)


def test_a_wobbling_domain_mask_does_not_repaint_the_boundary():
    """A one-pixel rim wobble is the mask being rebuilt, not the island moving.

    The measured case: on the shipped finger island this is a 0.73% area flip that
    repaints 69% of the boundary band and 35% of the rendered edge pixels.
    """
    latch = backend.DomainMaskLatch()
    image = _blank_image()
    # A zero field, so nothing but the boundary can change the picture.
    flow = np.zeros((*DOMAIN_HW, 2), dtype=np.float32)
    depth = np.ones(DOMAIN_HW, dtype=np.float32)
    wobbles = [(0, 0), (1, 0), (0, -1), (1, 1), (-1, 0), (0, 1), (-1, -1)]
    arrows, maps = [], []
    for dy, dx in wobbles:
        mask = _shift(_island(), dy, dx)
        frame = backend.ImageFrame(image, flow=flow, depth=depth, valid_mask=mask)
        used = latch.update(mask)
        arrows.append(
            backend.render_image(
                frame, mode="arrows_pure", draw_style="finger", domain_mask=used
            )
        )
        maps.append(
            backend.render_image(
                frame, mode="depth", depth_range=(0.0, 1.0), domain_mask=used
            )
        )
    assert all(np.array_equal(arrows[0], other) for other in arrows), (
        "the arrow canvas repainted a boundary that had not moved"
    )
    assert all(np.array_equal(maps[0], other) for other in maps), (
        "the depth view's no-data boundary repainted a boundary that had not moved"
    )
    assert latch.adoptions == 1, "a rim wobble was taken for a moved boundary"


def test_a_real_move_of_the_domain_is_followed():
    """The latch must not freeze a boundary that actually moved."""
    latch = backend.DomainMaskLatch()
    island = _island()
    latch.update(island)
    moved = _shift(island, 6, 0)
    used = latch.update(moved)
    assert np.array_equal(used, moved), "a moved island kept the old boundary"
    assert latch.adoptions == 2

    frame = backend.ImageFrame(
        _blank_image(), flow=np.zeros((*DOMAIN_HW, 2), np.float32)
    )
    before = backend.render_image(
        frame, mode="arrows_pure", draw_style="finger", domain_mask=island
    )
    after = backend.render_image(
        frame, mode="arrows_pure", draw_style="finger", domain_mask=used
    )
    assert not np.array_equal(before, after)


def test_a_frame_without_a_mask_cannot_move_the_boundary():
    latch = backend.DomainMaskLatch()
    assert latch.update(None) is None, "a session with no mask invented one"
    island = _island()
    latch.update(island)
    assert np.array_equal(latch.update(None), island), (
        "a frame that states no mask dropped the session's boundary"
    )
    assert (latch.adoptions, latch.held) == (1, 0)


def test_the_latched_boundary_never_gates_the_arrows():
    """Latching the boundary must not put an arrow on an unvouched-for sample."""
    image = _blank_image()
    flow = np.zeros((*DOMAIN_HW, 2), dtype=np.float32)
    flow[::10, ::10, 0] = 3.0
    # The frame claims nothing; the session's boundary is the whole domain.
    frame = backend.ImageFrame(image, flow=flow, valid_mask=np.zeros(DOMAIN_HW, bool))
    boundary = np.ones(DOMAIN_HW, dtype=bool)
    got = backend.render_image(
        frame, mode="arrows_pure", draw_style="finger", domain_mask=boundary
    )
    expected = draw_arrow_display(
        image,
        flow,
        background="pure_color",
        draw_style="finger",
        valid_mask=frame.valid_mask,
        domain_mask=boundary,
        colormap=flow_colormap("dark"),
    )
    assert np.array_equal(got, expected)
    # A whole-domain boundary paints one tone; an empty frame mask means no arrows,
    # so the picture stays that one tone rather than gaining arrow pixels.
    tones = {tuple(int(v) for v in tone) for tone in np.unique(got.reshape(-1, 3), axis=0)}
    assert tones == {(0x16, 0x13, 0x12)}, (
        f"arrows were drawn from a mask the frame does not claim: {tones}"
    )


def test_resetting_the_latch_clears_the_session_boundary():
    latch = backend.DomainMaskLatch()
    latch.update(_island())
    latch.reset()
    assert latch.mask is None
    assert (latch.adoptions, latch.held) == (0, 0)
    other = _shift(_island(), 5, 5)
    assert np.array_equal(latch.update(other), other), (
        "a new session kept the previous session's island"
    )