"""
test_pointing.py — Unit tests for pointing ray helpers (no ROS2 required).

Covers:
  - unproject: pixel + depth → 3-D point
  - sample_depth: median depth sampling with outlier rejection
  - compute_pointing_ray: full 3-D ray computation
"""

import sys
import os
import math

import numpy as np
import pytest

sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))

from saltybot_social._pointing_ray import unproject, sample_depth, compute_pointing_ray


# ── unproject ─────────────────────────────────────────────────────────────────

class TestUnproject:

    def test_principal_point_gives_zero_xy(self):
        X, Y, Z = unproject(320.0, 240.0, 1.0, fx=500.0, fy=500.0, cx=320.0, cy=240.0)
        assert X == pytest.approx(0.0)
        assert Y == pytest.approx(0.0)
        assert Z == pytest.approx(1.0)

    def test_z_equals_depth(self):
        _, _, Z = unproject(0.0, 0.0, 3.5, fx=600.0, fy=600.0, cx=320.0, cy=240.0)
        assert Z == pytest.approx(3.5)

    def test_right_of_principal_gives_positive_x(self):
        X, _, _ = unproject(420.0, 240.0, 1.0, fx=500.0, fy=500.0, cx=320.0, cy=240.0)
        assert X > 0.0

    def test_below_principal_gives_positive_y(self):
        _, Y, _ = unproject(320.0, 340.0, 1.0, fx=500.0, fy=500.0, cx=320.0, cy=240.0)
        assert Y > 0.0

    def test_x_scales_linearly_with_depth(self):
        X1, _, _ = unproject(400.0, 240.0, 1.0, fx=500.0, fy=500.0, cx=320.0, cy=240.0)
        X2, _, _ = unproject(400.0, 240.0, 2.0, fx=500.0, fy=500.0, cx=320.0, cy=240.0)
        assert X2 == pytest.approx(X1 * 2.0)

    def test_known_value(self):
        # (u=370, v=290, d=2.0, fx=fy=500, cx=320, cy=240)
        # X = (370-320)*2/500 = 0.2, Y = (290-240)*2/500 = 0.2, Z = 2.0
        X, Y, Z = unproject(370.0, 290.0, 2.0, fx=500.0, fy=500.0, cx=320.0, cy=240.0)
        assert X == pytest.approx(0.2)
        assert Y == pytest.approx(0.2)
        assert Z == pytest.approx(2.0)


# ── sample_depth ──────────────────────────────────────────────────────────────

class TestSampleDepth:

    def _uniform_depth(self, h=480, w=640, val=2.0):
        return np.full((h, w), val, dtype=np.float32)

    def test_uniform_image_returns_value(self):
        img = self._uniform_depth(val=1.5)
        d = sample_depth(img, 320.0, 240.0)
        assert d == pytest.approx(1.5)

    def test_out_of_range_returns_nan(self):
        img = np.zeros((480, 640), dtype=np.float32)   # all zeros (< d_min=0.1)
        d = sample_depth(img, 320.0, 240.0)
        assert math.isnan(d)

    def test_ignores_outliers_above_d_max(self):
        img = self._uniform_depth(val=1.0)
        # Put a large spike at the centre
        img[240, 320] = 100.0
        d = sample_depth(img, 320.0, 240.0, window=3, d_max=8.0)
        assert d == pytest.approx(1.0)

    def test_edge_pixel_does_not_crash(self):
        img = self._uniform_depth(val=2.0)
        d = sample_depth(img, 0.0, 0.0)
        assert d == pytest.approx(2.0)

    def test_median_of_mixed_values(self):
        img = np.zeros((10, 10), dtype=np.float32)
        img[5, 5] = 1.0
        img[5, 6] = 2.0
        img[6, 5] = 3.0
        d = sample_depth(img, 5.0, 5.0, window=3, d_min=0.0)
        # Valid: 1.0, 2.0, 3.0 — median = 2.0
        assert d == pytest.approx(2.0)


# ── compute_pointing_ray ──────────────────────────────────────────────────────

def _make_landmarks(mcp_nx=0.5, mcp_ny=0.5, tip_nx=0.6, tip_ny=0.4):
    """21 dummy landmarks; only MCP (idx 5) and TIP (idx 8) matter."""
    lms = [(0.5, 0.5, 0.0)] * 21
    lms[5] = (mcp_nx, mcp_ny, 0.0)   # INDEX_MCP
    lms[8] = (tip_nx, tip_ny, 0.0)   # INDEX_TIP
    return lms


def _depth_img(val=2.0, h=480, w=640):
    return np.full((h, w), val, dtype=np.float32)


_FX = 600.0; _FY = 600.0; _CX = 320.0; _CY = 240.0
_W  = 640;   _H  = 480


class TestComputePointingRay:

    def test_returns_dict_on_valid_input(self):
        lms = _make_landmarks()
        result = compute_pointing_ray(
            lms, _depth_img(), _FX, _FY, _CX, _CY, _W, _H)
        assert result is not None
        assert 'origin_3d' in result
        assert 'direction_3d' in result
        assert 'target_3d' in result
        assert 'range_m' in result

    def test_returns_none_when_depth_all_zero(self):
        lms = _make_landmarks()
        result = compute_pointing_ray(
            lms, np.zeros((_H, _W), dtype=np.float32),
            _FX, _FY, _CX, _CY, _W, _H)
        assert result is None

    def test_direction_is_unit_vector(self):
        lms = _make_landmarks(mcp_nx=0.4, mcp_ny=0.5, tip_nx=0.6, tip_ny=0.4)
        result = compute_pointing_ray(
            lms, _depth_img(), _FX, _FY, _CX, _CY, _W, _H)
        assert result is not None
        dx, dy, dz = result['direction_3d']
        norm = math.sqrt(dx*dx + dy*dy + dz*dz)
        assert norm == pytest.approx(1.0, abs=1e-5)

    def test_target_is_origin_plus_range_times_direction(self):
        lms = _make_landmarks()
        result = compute_pointing_ray(
            lms, _depth_img(val=2.0), _FX, _FY, _CX, _CY, _W, _H,
            ref_distance=3.0)
        ox, oy, oz = result['origin_3d']
        dx, dy, dz = result['direction_3d']
        tx, ty, tz = result['target_3d']
        r = result['range_m']
        assert tx == pytest.approx(ox + dx * r, abs=1e-5)
        assert ty == pytest.approx(oy + dy * r, abs=1e-5)
        assert tz == pytest.approx(oz + dz * r, abs=1e-5)

    def test_mcp_uv_matches_landmark_projection(self):
        lms = _make_landmarks(mcp_nx=0.5, mcp_ny=0.5)
        result = compute_pointing_ray(
            lms, _depth_img(), _FX, _FY, _CX, _CY, _W, _H)
        mu, mv = result['mcp_uv']
        assert mu == pytest.approx(0.5 * _W)
        assert mv == pytest.approx(0.5 * _H)

    def test_fallback_when_only_one_depth_valid(self):
        """Only the MCP pixel has valid depth; TIP is in an invalid region."""
        depth = np.zeros((_H, _W), dtype=np.float32)
        # Set valid depth only at the MCP location
        mcp_u, mcp_v = int(0.5 * _W), int(0.5 * _H)
        for du in range(-3, 4):
            for dv in range(-3, 4):
                u_ = max(0, min(_W - 1, mcp_u + du))
                v_ = max(0, min(_H - 1, mcp_v + dv))
                depth[v_, u_] = 2.0

        lms = _make_landmarks(mcp_nx=0.5, mcp_ny=0.5, tip_nx=0.9, tip_ny=0.1)
        result = compute_pointing_ray(
            lms, depth, _FX, _FY, _CX, _CY, _W, _H)
        # Should still return a result (using MCP depth for TIP fallback)
        assert result is not None
        dx, dy, dz = result['direction_3d']
        norm = math.sqrt(dx*dx + dy*dy + dz*dz)
        assert norm == pytest.approx(1.0, abs=1e-4)

    def test_ref_distance_used_as_range(self):
        lms = _make_landmarks()
        result = compute_pointing_ray(
            lms, _depth_img(), _FX, _FY, _CX, _CY, _W, _H,
            ref_distance=5.0)
        assert result['range_m'] == pytest.approx(5.0)


if __name__ == '__main__':
    pytest.main([__file__, '-v'])
