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