"""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_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"
    )