"""Shared 2D/3D flow sampling and domain-mask rules."""

from __future__ import annotations

import numpy as np
import pytest
from orisys_app_shared.scene3d.contracts import ThreeDError
from orisys_app_shared.scene3d.visual_policy import (
    VisualPolicy,
    build_domain_mask,
    paint_domain_background,
    select_flow_samples,
)


def test_build_domain_mask_rejects_shape_mismatch() -> None:
    geo = np.ones((4, 5), dtype=bool)
    with pytest.raises(ThreeDError):
        build_domain_mask(geo, shape_hw=(3, 5))


def test_build_domain_mask_intersects_roi() -> None:
    geo = np.ones((3, 3), dtype=bool)
    roi = np.zeros((3, 3), dtype=bool)
    roi[1, 1] = True
    mask, diags = build_domain_mask(geo, roi)
    assert mask is not None
    assert bool(mask[1, 1]) and not bool(mask[0, 0])
    assert diags == ()


def test_select_flow_samples_uses_shared_grid_and_mask() -> None:
    h, w = 40, 40
    flow = np.zeros((h, w, 2), dtype=np.float64)
    flow[15, 15] = (0.5, 0.0)  # |2*flow| = 1.0
    mask = np.ones((h, w), dtype=bool)
    mask[:, :5] = False
    policy = VisualPolicy(flow_on_px=0.4, flow_grid_spacing=10, flow_edge_margin=0)
    samples = select_flow_samples(flow, mask, policy)
    assert samples.rows.size >= 1
    assert all(c >= 5 for c in samples.cols.tolist())
    pairs = set(zip(samples.rows.tolist(), samples.cols.tolist()))
    assert (15, 15) in pairs


def test_select_flow_samples_hysteresis_keeps_weaker_motion() -> None:
    """The band between ``flow_off_px`` and ``flow_on_px`` decides through ``active``.

    The location is a candidate either way — the band is a statement about activity, not
    about where the policy samples — so what changes with ``previously_on`` is the flag.
    This is the single place the two views' thresholds are applied: a caller that draws
    ``active_only()`` draws what it drew when the selector filtered for it.
    """
    h, w = 30, 30
    flow = np.zeros((h, w, 2), dtype=np.float64)
    # |2*flow| = 0.7: below on (0.8) but above off (0.5)
    flow[15, 15] = (0.35, 0.0)
    policy = VisualPolicy(
        flow_on_px=0.8,
        flow_off_px=0.5,
        flow_grid_spacing=10,
        flow_edge_margin=0,
    )
    previously = np.zeros((h, w), dtype=bool)
    previously[15, 15] = True
    off = select_flow_samples(flow, None, policy, previously_on=None)
    on = select_flow_samples(flow, None, policy, previously_on=previously)
    # Same candidate locations: the grid does not depend on the verdict.
    assert off.rows.tolist() == on.rows.tolist()
    assert off.cols.tolist() == on.cols.tolist()
    assert (15, 15) in set(zip(on.rows.tolist(), on.cols.tolist()))
    # Only the verdict does.
    assert on.active_count == 1
    assert off.active_count == 0
    assert on.active_only().rows.tolist() == [15]
    assert on.active_only().cols.tolist() == [15]
    assert off.active_only().rows.size == 0


def test_the_location_cap_is_a_location_budget() -> None:
    """``max_arrows`` limits candidate locations, not arrows.

    The cap has to be spent before the activity verdict, otherwise the set of locations
    would depend on the threshold and a caller drawing inactive candidates too would draw
    a different set from the one a caller drawing arrows does. The strongest candidates
    are kept, and every inactive candidate is one no arrow was ever drawn for.
    """
    h, w = 40, 40
    flow = np.zeros((h, w, 2), dtype=np.float64)
    flow[15, 15] = (5.0, 0.0)  # far above the threshold
    policy = VisualPolicy(flow_on_px=0.4, flow_grid_spacing=1, flow_edge_margin=0, max_arrows=4)
    samples = select_flow_samples(flow, np.ones((h, w), dtype=bool), policy)
    assert samples.rows.size == 4
    assert samples.active_count == 1
    assert (15, 15) in set(zip(samples.rows.tolist(), samples.cols.tolist())), (
        "the cap dropped the only moving sample"
    )


def test_the_threshold_decides_the_marker_and_never_the_location() -> None:
    """The two markers are one grid read two ways: only the verdict moves with a threshold.

    The 2D canvas draws an arrow for an ``active`` sample and a quiet dot for an inactive
    one, both at the candidate location this selector returns. So the candidate set — the
    marker set, once each sample carries exactly one marker — is the policy's grid and
    nothing else, while ``flow_on_px`` decides how many of them are arrows. A threshold
    that moved a location would move the idle grid with it, which is what makes the dots
    comparable from frame to frame.
    """
    h, w = 30, 30
    flow = np.zeros((h, w, 2), dtype=np.float64)
    # 6 px displacements of 0.1 .. 1.5 px, each on the 5 px grid (2 + 5k) the policy
    # samples at, so each threshold step moves exactly one sample across.
    for index, value in enumerate((0.05, 0.1, 0.2, 0.35, 0.5, 0.75)):
        flow[2 + 5 * index, 2 + 5 * index] = (value, 0.0)

    locations: list[tuple[int, int]] | None = None
    arrows: list[int] = []
    for on_px in (0.2, 0.4, 0.6, 0.9, 1.2):
        policy = VisualPolicy(
            flow_on_px=on_px,
            flow_off_px=max(0.0, on_px - 0.1),
            flow_grid_spacing=5,
            flow_edge_margin=0,
        ).normalised()
        samples = select_flow_samples(flow, None, policy)
        here = list(zip(samples.rows.tolist(), samples.cols.tolist()))
        if locations is None:
            locations = here
        assert here == locations, (
            f"the candidate locations moved with the threshold at {on_px} px"
        )
        # One marker per sample, and which marker is the threshold's whole decision.
        assert samples.active.size == len(here)
        assert samples.active_count + int((~samples.active).sum()) == len(here)
        arrows.append(samples.active_count)
    assert arrows == [5, 4, 3, 2, 1], f"the sweep did not move one sample per step: {arrows}"


def test_paint_domain_background_colours() -> None:
    mask = np.zeros((4, 4), dtype=bool)
    mask[1:3, 1:3] = True
    canvas = paint_domain_background(
        (4, 4), mask, island_bgr=(10, 20, 30), exterior_bgr=(1, 2, 3)
    )
    assert tuple(canvas[0, 0]) == (1, 2, 3)
    assert tuple(canvas[1, 1]) == (10, 20, 30)