"""
Unit tests for saltybot_uwb_calibration.mds_math (Issue #602).
No ROS2 or hardware required — pure numpy.
"""

import math
import sys
import os

import numpy as np
import pytest

sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
from saltybot_uwb_calibration.mds_math import mds_2d, anchor_frame_align, _residual


def _dist_matrix(positions: np.ndarray) -> np.ndarray:
    """Build exact pairwise distance matrix from ground-truth positions."""
    n = len(positions)
    D = np.zeros((n, n))
    for i in range(n):
        for j in range(n):
            D[i, j] = np.linalg.norm(positions[i] - positions[j])
    return D


# ── Square configuration (4 anchors) ──────────────────────────────────────

GT_SQUARE = np.array([
    [0.0, 0.0],
    [4.0, 0.0],
    [4.0, 3.0],
    [0.0, 3.0],
])


class TestMdsSquare:
    def test_residual_near_zero(self):
        D = _dist_matrix(GT_SQUARE)
        positions, rms = mds_2d(D)
        assert rms < 1e-6, f"Residual too large: {rms}"

    def test_correct_distances_preserved(self):
        D = _dist_matrix(GT_SQUARE)
        positions, _ = mds_2d(D)
        positions = anchor_frame_align(positions)
        D_rec = _dist_matrix(positions)
        # All pairwise distances should match to < 1 mm
        assert np.max(np.abs(D - D_rec)) < 1e-3

    def test_anchor_0_at_origin(self):
        D = _dist_matrix(GT_SQUARE)
        positions, _ = mds_2d(D)
        positions = anchor_frame_align(positions)
        assert abs(positions[0, 0]) < 1e-9
        assert abs(positions[0, 1]) < 1e-9

    def test_anchor_1_on_positive_x(self):
        D = _dist_matrix(GT_SQUARE)
        positions, _ = mds_2d(D)
        positions = anchor_frame_align(positions)
        assert positions[1, 0] > 0
        assert abs(positions[1, 1]) < 1e-9


# ── Non-uniform configuration (5 anchors) ─────────────────────────────────

GT_5 = np.array([
    [0.0, 0.0],
    [5.0, 0.0],
    [5.0, 4.0],
    [2.5, 6.0],
    [0.0, 4.0],
])


class TestMds5Anchors:
    def test_residual_near_zero(self):
        D = _dist_matrix(GT_5)
        positions, rms = mds_2d(D)
        assert rms < 1e-6

    def test_shape(self):
        D = _dist_matrix(GT_5)
        positions, _ = mds_2d(D)
        assert positions.shape == (5, 2)


# ── Noisy distances (simulate UWB measurement noise) ──────────────────────

class TestMdsNoisyInput:
    def test_residual_acceptable_with_5cm_noise(self):
        rng = np.random.default_rng(42)
        D = _dist_matrix(GT_SQUARE)
        noise = rng.normal(0, 0.05, D.shape)
        noise = (noise + noise.T) / 2   # keep symmetric
        np.fill_diagonal(noise, 0)
        D_noisy = np.maximum(D + noise, 0)
        positions, rms = mds_2d(D_noisy)
        # With 5 cm noise the residual should be sub-10 cm
        assert rms < 0.10, f"Residual {rms:.4f} m too large"


# ── Error handling ─────────────────────────────────────────────────────────

class TestMdsErrors:
    def test_non_square_raises(self):
        with pytest.raises(ValueError):
            mds_2d(np.ones((3, 4)))

    def test_too_few_anchors_raises(self):
        with pytest.raises(ValueError):
            mds_2d(np.array([[0.0, 1.0], [1.0, 0.0]]))

    def test_asymmetric_raises(self):
        D = _dist_matrix(GT_SQUARE)
        D[0, 1] += 0.5  # break symmetry
        with pytest.raises(ValueError):
            mds_2d(D)


# ── anchor_frame_align ─────────────────────────────────────────────────────

class TestFrameAlign:
    def test_translated_positions_aligned(self):
        pts = GT_SQUARE.copy() + np.array([10.0, -5.0])
        aligned = anchor_frame_align(pts)
        assert abs(aligned[0, 0]) < 1e-9
        assert abs(aligned[0, 1]) < 1e-9

    def test_rotated_positions_aligned(self):
        angle = math.pi / 3
        R = np.array([[math.cos(angle), -math.sin(angle)],
                      [math.sin(angle),  math.cos(angle)]])
        pts = (R @ GT_SQUARE.T).T
        aligned = anchor_frame_align(pts)
        assert aligned[1, 0] > 0
        assert abs(aligned[1, 1]) < 1e-9


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