Presa-Mano/tests/test_work_plane.py

"""Unit tests for perception.work_plane."""
from __future__ import annotations

import json

import numpy as np
import pytest

from presa_mano.perception.work_plane import (
    WorkPlane,
    WorkPlaneError,
    apply_below_plane_filter,
    fit_plane_svd,
    _load_work_plane_from_disk,
    _save_work_plane_to_disk,
    _work_plane_path,
)


def _flat_plane_cloud(z: float = 1.0, n: int = 200, noise_m: float = 0.0):
    """Cloud sampled on the plane z=z (camera frame, OpenCV: camera looks +Z)."""
    rng = np.random.default_rng(0)
    xy = rng.uniform(-0.2, 0.2, size=(n, 2))
    zs = np.full((n, 1), z) + rng.normal(0, noise_m, size=(n, 1))
    return np.hstack([xy, zs])


def test_fit_plane_svd_horizontal_plane():
    """Plane at z=1.0, normal must point toward camera (negative Z)."""
    pts = _flat_plane_cloud(z=1.0, n=200, noise_m=0.0)
    normal, offset, centroid, rms_mm = fit_plane_svd(pts)
    assert abs(normal[2]) > 0.99
    assert normal[2] < 0
    assert abs(offset - 1.0) < 1e-3
    assert rms_mm < 1.0


def test_fit_plane_svd_raises_on_too_few_points():
    with pytest.raises(WorkPlaneError):
        fit_plane_svd(np.zeros((5, 3)))


def test_apply_below_plane_filter_noop_without_plane():
    pts = _flat_plane_cloud(z=1.0)
    out, info = apply_below_plane_filter(pts, None)
    assert info["plane_applied"] is False
    assert info["n_points_before_plane"] == len(pts)
    assert out is pts


def test_apply_below_plane_filter_keeps_above_drops_below():
    """Plane at z=1.0, normal toward camera (-Z): points with z<1.0 are above."""
    above = _flat_plane_cloud(z=0.95, n=100)
    below = _flat_plane_cloud(z=1.10, n=100)
    pts = np.vstack([above, below])
    wp = WorkPlane(
        normal=[0.0, 0.0, -1.0], offset=1.0,
        centroid=[0.0, 0.0, 1.0],
        n_inliers=100, fit_rms_mm=0.0,
        roi=[0, 0, 10, 10],
        offset_mm=0.0,
        created_at=0.0,
    )
    out, info = apply_below_plane_filter(pts, wp)
    assert info["plane_applied"] is True
    assert info["n_points_before_plane"] == 200
    assert len(out) == 100
    assert np.all(out[:, 2] <= 1.0 + 1e-9)


def test_apply_below_plane_filter_offset_buffer_keeps_below():
    """offset_mm=50 keeps points up to 50mm below the plane."""
    above = _flat_plane_cloud(z=0.95, n=50)
    just_below = _flat_plane_cloud(z=1.03, n=50)
    far_below = _flat_plane_cloud(z=1.20, n=50)
    pts = np.vstack([above, just_below, far_below])
    wp = WorkPlane(
        normal=[0.0, 0.0, -1.0], offset=1.0,
        centroid=[0.0, 0.0, 1.0],
        n_inliers=50, fit_rms_mm=0.0,
        roi=[0, 0, 10, 10],
        offset_mm=50.0,
        created_at=0.0,
    )
    out, info = apply_below_plane_filter(pts, wp)
    assert info["plane_applied"] is True
    assert len(out) == 100   # above (50) + just_below (50); far_below dropped


def test_apply_below_plane_filter_empty_input():
    wp = WorkPlane(
        normal=[0.0, 0.0, -1.0], offset=1.0,
        centroid=[0.0, 0.0, 1.0],
        n_inliers=10, fit_rms_mm=0.0,
        roi=[0, 0, 10, 10],
        offset_mm=5.0, created_at=0.0,
    )
    out, info = apply_below_plane_filter(np.empty((0, 3)), wp)
    assert info["plane_applied"] is False   # short-circuit: empty input never "ran"
    assert info["n_points_before_plane"] == 0
    assert len(out) == 0


def test_work_plane_disk_roundtrip(tmp_path, monkeypatch):
    """Save then load preserves all fields including offset_mm."""
    monkeypatch.setenv("PRESA_MANO_DATA", str(tmp_path))
    wp = WorkPlane(
        normal=[0.0, 0.0, -1.0], offset=1.0,
        centroid=[0.1, 0.2, 1.0],
        n_inliers=42, fit_rms_mm=1.5,
        roi=[10, 20, 100, 200],
        offset_mm=7.5,
        created_at=1234567890.0,
    )
    _save_work_plane_to_disk(wp)
    loaded = _load_work_plane_from_disk()
    assert loaded is not None
    assert loaded.normal == wp.normal
    assert loaded.offset == wp.offset
    assert loaded.centroid == wp.centroid
    assert loaded.n_inliers == wp.n_inliers
    assert loaded.fit_rms_mm == wp.fit_rms_mm
    assert loaded.roi == wp.roi
    assert loaded.offset_mm == 7.5
    assert loaded.created_at == wp.created_at


def test_work_plane_load_legacy_file_defaults_offset(tmp_path, monkeypatch):
    """A pre-existing work_plane.json without offset_mm gets the 5.0 default."""
    monkeypatch.setenv("PRESA_MANO_DATA", str(tmp_path))
    path = _work_plane_path()
    legacy = {
        "normal": [0.0, 0.0, -1.0], "offset": 1.0,
        "centroid": [0.0, 0.0, 1.0],
        "n_inliers": 10, "fit_rms_mm": 1.0,
        "roi": [0, 0, 10, 10],
        "created_at": 0.0,
    }
    path.write_text(json.dumps(legacy))
    loaded = _load_work_plane_from_disk()
    assert loaded is not None
    assert loaded.offset_mm == 5.0