"""Composer overlays and 2D/3D sample-point parity."""

from __future__ import annotations

import numpy as np

from orisys_app_shared.acquisition.arrows import _policy_from_style
from orisys_app_shared.scene3d.composer import SceneComposer
from orisys_app_shared.scene3d.contracts import (
    DisplaySettings,
    FieldDescriptor,
    SurfaceFieldFrame,
    freeze,
)
from orisys_app_shared.scene3d.synthetic import poc_identity, poc_scene
from orisys_app_shared.scene3d.visual_policy import VisualPolicy, select_flow_samples


def _depth_frame(assets, *, flow=None, contact_xy=None, fn=0.0):
    h, w = assets.shape_hw
    scalar = freeze(np.zeros((h, w), dtype=np.float32), dtype=np.float32)
    valid = freeze(np.ones((h, w), dtype=bool), dtype=np.bool_)
    flow_uv = None if flow is None else freeze(np.asarray(flow, dtype=np.float32), dtype=np.float32)
    flow_mask = None if flow is None else freeze(np.asarray(valid, dtype=bool).copy(), dtype=np.bool_)
    found = contact_xy is not None
    return SurfaceFieldFrame(
        identity=poc_identity(),
        geometry_id=assets.geometry_id,
        field_domain_id=assets.field_domain_id,
        scalar=scalar,
        valid=valid,
        scalar_descriptor=FieldDescriptor(
            "depth_response", "1", "approximate", False
        ),
        computed_frame_id=0,
        flow_uv=flow_uv,
        flow_valid_mask=flow_mask,
        contact_uv_xy=None if contact_xy is None else tuple(contact_xy),
        contact_found=found,
        contact_valid=found,
        contact_fn=float(fn),
    )


def test_composer_builds_overlays_when_mapping_present() -> None:
    assets = poc_scene()
    h, w = assets.shape_hw
    flow = np.zeros((h, w, 2), dtype=np.float32)
    flow[1, 1] = (0.9, 0.0)
    frame = _depth_frame(assets, flow=flow, contact_xy=(1.0, 1.0), fn=800.0)
    composed = SceneComposer().compose(
        assets,
        frame,
        DisplaySettings(
            flow_arrows_enabled=True,
            contact_arrow_enabled=True,
            flow_grid_spacing=1,
            flow_on_px=0.4,
            flow_off_px=0.3,
            max_arrows=50,
        ),
    )
    assert composed.flow_arrows is not None
    assert composed.flow_arrows.origins.shape[0] >= 1
    assert composed.contact_arrow is not None
    assert composed.contact_arrow.origins.shape[0] == 1


def test_composer_hides_overlays_when_disabled() -> None:
    assets = poc_scene()
    h, w = assets.shape_hw
    flow = np.zeros((h, w, 2), dtype=np.float32)
    flow[1, 1] = (0.9, 0.0)
    frame = _depth_frame(assets, flow=flow, contact_xy=(1.0, 1.0), fn=800.0)
    composed = SceneComposer().compose(
        assets,
        frame,
        DisplaySettings(flow_arrows_enabled=False, contact_arrow_enabled=False),
    )
    assert composed.flow_arrows is None
    assert composed.contact_arrow is None


def test_2d_3d_flow_sample_points_match() -> None:
    assets = poc_scene()
    h, w = assets.shape_hw
    flow = np.zeros((h, w, 2), dtype=np.float64)
    flow[1, 1] = (0.7, 0.2)
    mask = np.ones((h, w), dtype=bool)
    settings = DisplaySettings(
        flow_on_px=0.4,
        flow_off_px=0.3,
        flow_grid_spacing=1,
        max_arrows=50,
    )
    policy_3d = _policy_from_style(
        "finger",
        flow_on_px=settings.flow_on_px,
        flow_off_px=settings.flow_off_px,
        flow_grid_spacing=settings.flow_grid_spacing,
        max_arrows=settings.max_arrows,
    )
    policy_shared = VisualPolicy(
        flow_on_px=settings.flow_on_px,
        flow_off_px=settings.flow_off_px,
        flow_grid_spacing=settings.flow_grid_spacing,
        max_arrows=settings.max_arrows,
        flow_edge_margin=0,
    ).normalised()
    samples_2d = select_flow_samples(flow, mask, policy_shared)
    samples_3d = select_flow_samples(flow, mask, policy_shared)
    assert list(zip(samples_2d.rows, samples_2d.cols)) == list(
        zip(samples_3d.rows, samples_3d.cols)
    )
    assert policy_3d.flow_on_px == settings.flow_on_px