"""
Unit tests for person_follower_node logic.
No ROS2 runtime required — tests pure functions mirrored from the node.
Run with: pytest jetson/ros2_ws/src/saltybot_follower/test/test_person_follower.py
"""

import math
import pytest


# ── Pure function mirrors (no ROS2 imports) ──────────────────────────────────

def _clamp(v, lo, hi):
    return max(lo, min(hi, v))


def _compute_distance(px, py):
    return math.sqrt(px * px + py * py)


def _compute_bearing(px, py):
    """Bearing to (px,py) in base_link: positive = person to the left."""
    return math.atan2(py, px)


def _compute_follow_cmd(px, py,
                        follow_distance, dead_zone,
                        kp_linear, kp_angular,
                        max_linear_vel, max_angular_vel,
                        obstacle_ahead=False):
    """Mirror of PersonFollowerNode proportional controller."""
    distance = _compute_distance(px, py)
    bearing  = _compute_bearing(px, py)

    dist_error = distance - follow_distance
    if abs(dist_error) > dead_zone:
        linear_x = _clamp(kp_linear * dist_error, -max_linear_vel, max_linear_vel)
    else:
        linear_x = 0.0

    if obstacle_ahead and linear_x > 0.0:
        linear_x = 0.0

    angular_z = _clamp(kp_angular * bearing, -max_angular_vel, max_angular_vel)
    return linear_x, angular_z


def _check_obstacle(costmap_data, width, height, resolution,
                    check_dist, robot_half_width, threshold):
    """Mirror of check_obstacle_in_costmap."""
    if resolution <= 0.0:
        return False
    cx = width  // 2
    cy = height // 2
    n_fwd  = int(check_dist / resolution)
    n_side = int(robot_half_width / resolution)
    for dx in range(1, n_fwd + 1):
        col = cx + dx
        for dy in range(-n_side, n_side + 1):
            row = cy + dy
            if row < 0 or row >= height or col < 0 or col >= width:
                continue
            idx = row * width + col
            if idx < len(costmap_data) and costmap_data[idx] >= threshold:
                return True
    return False


# ── Costmap helpers ───────────────────────────────────────────────────────────

def _empty_costmap(width=20, height=20):
    return [0] * (width * height)


def _costmap_with_obstacle(col, row, width=20, height=20, value=99):
    data = [0] * (width * height)
    data[row * width + col] = value
    return data


# ── Default params ────────────────────────────────────────────────────────────

_DEF = dict(
    follow_distance=1.5,
    dead_zone=0.3,
    kp_linear=0.6,
    kp_angular=1.2,
    max_linear_vel=0.5,
    max_angular_vel=1.0,
)


# ── Clamp tests ───────────────────────────────────────────────────────────────

class TestClamp:
    def test_within_range(self):
        assert _clamp(0.5, 0.0, 1.0) == 0.5

    def test_clamp_high(self):
        assert _clamp(2.0, 0.0, 1.0) == 1.0

    def test_clamp_low(self):
        assert _clamp(-2.0, -1.0, 0.0) == -1.0

    def test_at_boundary_high(self):
        assert _clamp(1.0, 0.0, 1.0) == 1.0

    def test_at_boundary_low(self):
        assert _clamp(0.0, 0.0, 1.0) == 0.0


# ── Geometry tests ────────────────────────────────────────────────────────────

class TestGeometry:
    def test_distance_straight_ahead(self):
        assert _compute_distance(2.0, 0.0) == pytest.approx(2.0)

    def test_distance_diagonal(self):
        assert _compute_distance(3.0, 4.0) == pytest.approx(5.0)

    def test_distance_behind(self):
        assert _compute_distance(-2.0, 0.0) == pytest.approx(2.0)

    def test_bearing_straight_ahead(self):
        assert _compute_bearing(1.0, 0.0) == pytest.approx(0.0)

    def test_bearing_to_the_left(self):
        # +y = left → bearing = +π/2
        assert _compute_bearing(0.0, 1.0) == pytest.approx(math.pi / 2)

    def test_bearing_to_the_right(self):
        # -y = right → bearing = -π/2
        assert _compute_bearing(0.0, -1.0) == pytest.approx(-math.pi / 2)

    def test_bearing_behind_left(self):
        # (-x, +y) → ≈ +3π/4
        assert _compute_bearing(-1.0, 1.0) == pytest.approx(3 * math.pi / 4)


# ── Follow command — linear axis ──────────────────────────────────────────────

class TestLinearCmd:
    def test_at_exact_target_zero_linear(self):
        lin, _ = _compute_follow_cmd(1.5, 0.0, **_DEF)
        assert lin == pytest.approx(0.0)

    def test_within_dead_zone_zero_linear(self):
        # error = 1.6 - 1.5 = 0.1 < dead_zone=0.3
        lin, _ = _compute_follow_cmd(1.6, 0.0, **_DEF)
        assert lin == pytest.approx(0.0)

    def test_outside_dead_zone_approach(self):
        # error = 2.0 - 1.5 = 0.5 → 0.6 * 0.5 = 0.3 (below max_linear_vel=0.5)
        lin, _ = _compute_follow_cmd(2.0, 0.0, **_DEF)
        assert lin == pytest.approx(0.3)

    def test_outside_dead_zone_retreat(self):
        # error = 1.0 - 1.5 = -0.5 → 0.6 * -0.5 = -0.3 (below max_linear_vel=0.5)
        lin, _ = _compute_follow_cmd(1.0, 0.0, **_DEF)
        assert lin == pytest.approx(-0.3)

    def test_linear_clamped_at_max(self):
        # error = 3.5 - 1.5 = 2.0 → 0.6 * 2.0 = 1.2 → clamped to 0.5
        lin, _ = _compute_follow_cmd(3.5, 0.0, **_DEF)
        assert lin == pytest.approx(0.5)

    def test_linear_clamped_at_min(self):
        # error = 0.0 - 1.5 = -1.5 → 0.6 * -1.5 = -0.9 → clamped to -0.5
        lin, _ = _compute_follow_cmd(0.0, 0.0, **_DEF)
        assert lin == pytest.approx(-0.5)

    def test_dead_zone_boundary_just_outside(self):
        # error = 2.5 - 1.5 + 0.3 + 0.01 = 0.31 → should produce output
        lin, _ = _compute_follow_cmd(1.5 + 0.3 + 0.01, 0.0, **_DEF)
        assert lin > 0.0

    def test_dead_zone_boundary_just_inside(self):
        lin, _ = _compute_follow_cmd(1.5 + 0.29, 0.0, **_DEF)
        assert lin == pytest.approx(0.0)


# ── Follow command — angular axis ─────────────────────────────────────────────

class TestAngularCmd:
    def test_person_ahead_zero_angular(self):
        _, ang = _compute_follow_cmd(2.0, 0.0, **_DEF)
        assert ang == pytest.approx(0.0)

    def test_angular_proportional_to_bearing(self):
        _, ang = _compute_follow_cmd(1.5, 0.5, **_DEF)
        expected = _clamp(1.2 * _compute_bearing(1.5, 0.5), -1.0, 1.0)
        assert ang == pytest.approx(expected, abs=1e-4)

    def test_angular_clamped_at_max(self):
        # Person far left: bearing ≈ π/2 → 1.2*π/2 >> 1.0 → clamped
        _, ang = _compute_follow_cmd(2.0, 10.0, **_DEF)
        assert ang == pytest.approx(1.0)

    def test_angular_clamped_at_min(self):
        _, ang = _compute_follow_cmd(2.0, -10.0, **_DEF)
        assert ang == pytest.approx(-1.0)

    def test_angular_corrects_within_dead_zone(self):
        # Even at target distance (linear=0), bearing should still be corrected
        lin, ang = _compute_follow_cmd(1.5, 0.5, **_DEF)
        assert lin == pytest.approx(0.0)   # dead zone
        assert ang != pytest.approx(0.0)   # bearing correction active

    def test_angular_corrects_when_retreating(self):
        # Robot too close and person is off to the side — should turn AND reverse
        lin, ang = _compute_follow_cmd(0.5, 0.5, **_DEF)
        assert lin < 0.0    # retreating
        assert ang != pytest.approx(0.0)


# ── Obstacle override tests ───────────────────────────────────────────────────

class TestObstacleOverride:
    def test_obstacle_suppresses_forward_motion(self):
        lin, _ = _compute_follow_cmd(2.5, 0.0, obstacle_ahead=True, **_DEF)
        assert lin == pytest.approx(0.0)

    def test_obstacle_does_not_suppress_reverse(self):
        # Too close (1.0m) → robot reverses; obstacle should NOT suppress reverse
        # error = 1.0 - 1.5 = -0.5 → kp * error = -0.3 (negative = reverse)
        lin, _ = _compute_follow_cmd(1.0, 0.0, obstacle_ahead=True, **_DEF)
        assert lin == pytest.approx(-0.3)

    def test_obstacle_allows_turning(self):
        lin, ang = _compute_follow_cmd(2.5, 2.0, obstacle_ahead=True, **_DEF)
        assert lin == pytest.approx(0.0)   # suppressed
        assert ang != pytest.approx(0.0)   # turning still active

    def test_no_obstacle_normal_forward(self):
        # error = 2.0 - 1.5 = 0.5 → 0.6 * 0.5 = 0.3
        lin, _ = _compute_follow_cmd(2.0, 0.0, obstacle_ahead=False, **_DEF)
        assert lin == pytest.approx(0.3)


# ── Costmap obstacle check tests ──────────────────────────────────────────────

class TestCostmapCheck:
    _RES = 0.05   # 5 cm/cell
    _W   = 20
    _H   = 20

    def _check(self, data, *, check_dist=0.5, half_width=0.15, threshold=70):
        return _check_obstacle(
            costmap_data=data,
            width=self._W,
            height=self._H,
            resolution=self._RES,
            check_dist=check_dist,
            robot_half_width=half_width,
            threshold=threshold,
        )

    def test_empty_costmap_no_obstacle(self):
        assert self._check(_empty_costmap(self._W, self._H)) is False

    def test_obstacle_one_cell_ahead_detected(self):
        # Robot centre col=10; dx=1 → col=11, row=10 (centreline)
        data = _costmap_with_obstacle(11, 10, self._W, self._H, value=99)
        assert self._check(data) is True

    def test_obstacle_at_check_dist_boundary_detected(self):
        # check_dist=0.5 / 0.05 = 10 cells → dx=1..10 → cols 11..20
        # Width=20 means valid cols 0..19; col 20 is out of bounds.
        # Last detectable col = cx + 9 = 10 + 9 = 19 (dx=9 within range(1,11))
        data = _costmap_with_obstacle(19, 10, self._W, self._H, value=99)
        assert self._check(data) is True

    def test_obstacle_behind_robot_not_detected(self):
        data = _costmap_with_obstacle(5, 10, self._W, self._H, value=99)
        assert self._check(data) is False

    def test_obstacle_within_corridor_side_detected(self):
        # half_width=0.15/0.05=3 cells; row=10+2=12 is inside corridor
        data = _costmap_with_obstacle(12, 12, self._W, self._H, value=99)
        assert self._check(data) is True

    def test_obstacle_outside_corridor_not_detected(self):
        # row=10+4=14 is outside 3-cell corridor
        data = _costmap_with_obstacle(12, 14, self._W, self._H, value=99)
        assert self._check(data) is False

    def test_value_below_threshold_not_detected(self):
        data = _costmap_with_obstacle(11, 10, self._W, self._H, value=50)
        assert self._check(data) is False

    def test_value_at_threshold_detected(self):
        data = _costmap_with_obstacle(11, 10, self._W, self._H, value=70)
        assert self._check(data) is True

    def test_zero_resolution_safe_returns_false(self):
        assert _check_obstacle(
            _empty_costmap(), width=20, height=20, resolution=0.0,
            check_dist=0.5, robot_half_width=0.15, threshold=70) is False

    def test_lethal_threshold_99_detected(self):
        data = _costmap_with_obstacle(11, 10, self._W, self._H, value=100)
        assert self._check(data, threshold=99) is True


# ── State machine tests ───────────────────────────────────────────────────────

class TestStateMachine:
    _FOLLOWING = "following"
    _STOPPING  = "stopping"
    _SEARCHING = "searching"

    def _transition(self, state, dt_lost, lost_timeout=2.0, search_timeout=5.0):
        """Mirror of _control_cb state transition block."""
        if state == self._FOLLOWING:
            if dt_lost > search_timeout:
                return self._SEARCHING
            elif dt_lost > lost_timeout:
                return self._STOPPING
        return state

    def test_following_stays_within_lost_timeout(self):
        assert self._transition(self._FOLLOWING, dt_lost=1.0) == self._FOLLOWING

    def test_following_to_stopping_past_lost_timeout(self):
        assert self._transition(self._FOLLOWING, dt_lost=2.5) == self._STOPPING

    def test_following_to_searching_past_search_timeout(self):
        assert self._transition(self._FOLLOWING, dt_lost=6.0) == self._SEARCHING

    def test_stopping_does_not_transition(self):
        # Only _target_cb can move out of STOPPING
        assert self._transition(self._STOPPING, dt_lost=3.0) == self._STOPPING

    def test_searching_does_not_transition(self):
        assert self._transition(self._SEARCHING, dt_lost=10.0) == self._SEARCHING

    def test_boundary_exactly_at_lost_timeout(self):
        # dt_lost == 2.0 is NOT > 2.0 — still FOLLOWING
        assert self._transition(self._FOLLOWING, dt_lost=2.0) == self._FOLLOWING

    def test_boundary_exactly_at_search_timeout(self):
        # dt_lost=5.0: NOT > search_timeout=5.0, but IS > lost_timeout=2.0 → STOPPING
        assert self._transition(self._FOLLOWING, dt_lost=5.0) == self._STOPPING

    def test_before_first_detection_initial_state_is_stopping(self):
        # Initial state is STOPPING (not FOLLOWING) so no transition occurs
        assert self._transition(self._STOPPING, dt_lost=1e9) == self._STOPPING

    def test_new_target_resets_to_following(self):
        # _target_cb sets state = FOLLOWING unconditionally
        state = self._SEARCHING
        # simulate _target_cb
        state = self._FOLLOWING
        assert state == self._FOLLOWING


# ── Search and stop behaviour ─────────────────────────────────────────────────

class TestSearchBehaviour:
    def test_searching_publishes_angular_only(self):
        state = "searching"
        search_angular_vel = 0.3
        linear_x  = 0.0 if state == "searching" else 99.9
        angular_z = search_angular_vel if state == "searching" else 0.0
        assert linear_x  == pytest.approx(0.0)
        assert angular_z == pytest.approx(0.3)

    def test_stopping_publishes_zero(self):
        state = "stopping"
        linear_x  = 0.0
        angular_z = 0.0
        assert linear_x  == pytest.approx(0.0)
        assert angular_z == pytest.approx(0.0)


# ── Follow-enabled gate ───────────────────────────────────────────────────────

class TestFollowEnabled:
    def test_disabled_zeroes_output(self):
        follow_enabled = False
        lin, ang = (0.0, 0.0) if not follow_enabled else \
            _compute_follow_cmd(3.0, 0.0, **_DEF)
        assert lin == pytest.approx(0.0)
        assert ang == pytest.approx(0.0)

    def test_enabled_passes_through(self):
        follow_enabled = True
        lin, ang = (0.0, 0.0) if not follow_enabled else \
            _compute_follow_cmd(3.0, 0.0, **_DEF)
        assert lin != pytest.approx(0.0)


# ── fuse_targets — UWB + camera fusion ──────────────────────────────────────

def _fuse_targets(uwb_x, uwb_y, uwb_fresh, cam_x, cam_y, cam_fresh, uwb_weight):
    """Mirror of person_follower_node.fuse_targets (no ROS2)."""
    def clamp(v, lo, hi): return max(lo, min(hi, v))
    if uwb_fresh and cam_fresh:
        w = clamp(uwb_weight, 0.0, 1.0)
        return w * uwb_x + (1.0 - w) * cam_x, w * uwb_y + (1.0 - w) * cam_y, True
    if uwb_fresh:
        return uwb_x, uwb_y, True
    if cam_fresh:
        return cam_x, cam_y, True
    return 0.0, 0.0, False


class TestFuseTargets:

    def test_both_fresh_blends(self):
        """Both sources fresh → weighted blend."""
        fx, fy, ok = _fuse_targets(
            uwb_x=4.0, uwb_y=0.0, uwb_fresh=True,
            cam_x=2.0, cam_y=0.0, cam_fresh=True,
            uwb_weight=0.7,
        )
        assert ok is True
        assert fx == pytest.approx(0.7 * 4.0 + 0.3 * 2.0)
        assert fy == pytest.approx(0.0)

    def test_uwb_only(self):
        """Only UWB fresh → UWB position returned as-is."""
        fx, fy, ok = _fuse_targets(
            uwb_x=3.0, uwb_y=1.0, uwb_fresh=True,
            cam_x=0.0, cam_y=0.0, cam_fresh=False,
            uwb_weight=0.7,
        )
        assert ok is True
        assert fx == pytest.approx(3.0)
        assert fy == pytest.approx(1.0)

    def test_camera_only(self):
        """Only camera fresh → camera position returned as-is."""
        fx, fy, ok = _fuse_targets(
            uwb_x=0.0, uwb_y=0.0, uwb_fresh=False,
            cam_x=2.5, cam_y=-0.5, cam_fresh=True,
            uwb_weight=0.7,
        )
        assert ok is True
        assert fx == pytest.approx(2.5)
        assert fy == pytest.approx(-0.5)

    def test_neither_fresh(self):
        """Both stale → (0, 0, False)."""
        fx, fy, ok = _fuse_targets(
            uwb_x=3.0, uwb_y=1.0, uwb_fresh=False,
            cam_x=2.0, cam_y=0.0, cam_fresh=False,
            uwb_weight=0.7,
        )
        assert ok is False
        assert fx == pytest.approx(0.0)
        assert fy == pytest.approx(0.0)

    def test_uwb_weight_zero_gives_camera(self):
        """uwb_weight=0.0 → pure camera even when both fresh."""
        fx, fy, ok = _fuse_targets(
            uwb_x=4.0, uwb_y=2.0, uwb_fresh=True,
            cam_x=1.0, cam_y=0.5, cam_fresh=True,
            uwb_weight=0.0,
        )
        assert ok is True
        assert fx == pytest.approx(1.0)
        assert fy == pytest.approx(0.5)

    def test_uwb_weight_one_gives_uwb(self):
        """uwb_weight=1.0 → pure UWB even when both fresh."""
        fx, fy, ok = _fuse_targets(
            uwb_x=4.0, uwb_y=2.0, uwb_fresh=True,
            cam_x=1.0, cam_y=0.5, cam_fresh=True,
            uwb_weight=1.0,
        )
        assert ok is True
        assert fx == pytest.approx(4.0)
        assert fy == pytest.approx(2.0)

    def test_lateral_blend(self):
        """Lateral (Y) component is also blended correctly."""
        fx, fy, ok = _fuse_targets(
            uwb_x=2.0, uwb_y=1.0, uwb_fresh=True,
            cam_x=2.0, cam_y=-1.0, cam_fresh=True,
            uwb_weight=0.5,
        )
        assert ok is True
        assert fx == pytest.approx(2.0)
        assert fy == pytest.approx(0.0)  # 0.5*1.0 + 0.5*(-1.0) = 0
