"""Unit tests for partition_by_plane."""
from __future__ import annotations

import numpy as np
import pytest

from presa_mano.perception.work_plane import WorkPlane, partition_by_plane


def _wp(offset_mm: float = 0.0) -> WorkPlane:
    """Work surface at z=0.5m in OpenCV camera frame (camera at origin
    looking +Z). Normal=[0,0,-1] points toward the camera — same convention
    as fit_plane_svd. Plane equation: -z + 0.5 = 0."""
    return WorkPlane(
        normal=[0.0, 0.0, -1.0],
        offset=0.5,
        centroid=[0.0, 0.0, 0.5],
        n_inliers=100,
        fit_rms_mm=0.5,
        roi=[0, 0, 100, 100],
        offset_mm=offset_mm,
        created_at=0.0,
    )


def test_partition_no_plane_returns_all_accepted():
    pts = np.array([[0.0, 0.0, 1.0], [1.0, 2.0, 3.0]])
    accepted, excluded, info = partition_by_plane(pts, None)
    assert np.array_equal(accepted, pts)
    assert excluded.shape == (0, 3)
    assert info == {"plane_applied": False, "n_points_before_plane": 2}


def test_partition_empty_input():
    pts = np.empty((0, 3), dtype=np.float64)
    accepted, excluded, info = partition_by_plane(pts, _wp())
    assert accepted.shape == (0, 3)
    assert excluded.shape == (0, 3)
    assert info["plane_applied"] is False
    assert info["n_points_before_plane"] == 0


def test_partition_splits_by_offset():
    """Plane at z=0.5, normal=-Z, offset_mm=5.

    Camera-side ("above") is z<0.5. Filter keeps points up to offset_mm
    *below* the plane (i.e. z<=0.505) and excludes farther-than-offset
    points (z>0.505)."""
    pts = np.array([
        [0.0, 0.0, 0.40],   # 100mm above → accepted
        [0.0, 0.0, 0.50],   # on plane → accepted (boundary)
        [0.0, 0.0, 0.504],  # 4mm below → accepted (within offset 5mm)
        [0.0, 0.0, 0.510],  # 10mm below → excluded
        [0.0, 0.0, 1.000],  # 500mm below → excluded
    ])
    accepted, excluded, info = partition_by_plane(pts, _wp(offset_mm=5.0))
    assert accepted.shape[0] == 3
    assert excluded.shape[0] == 2
    assert info["plane_applied"] is True
    assert info["n_points_before_plane"] == 5
    # partition: nessuna sovrapposizione
    union = np.vstack([accepted, excluded])
    assert sorted(union[:, 2].tolist()) == sorted(pts[:, 2].tolist())


def test_partition_consistent_with_apply_below_plane_filter():
    """`accepted` di partition_by_plane deve essere bit-identico a quello
    ritornato da apply_below_plane_filter (l'API legacy non deve cambiare)."""
    from presa_mano.perception.work_plane import apply_below_plane_filter
    rng = np.random.default_rng(42)
    pts = rng.normal(scale=0.1, size=(200, 3))
    wp = _wp(offset_mm=10.0)
    accepted_new, _excluded, _info = partition_by_plane(pts, wp)
    accepted_legacy, info_legacy = apply_below_plane_filter(pts, wp)
    np.testing.assert_array_equal(accepted_new, accepted_legacy)
    assert info_legacy["n_points_before_plane"] == 200