"""
test_seg_utils.py — Unit tests for seg_utils pure helpers.

No ROS2 / TensorRT / cv2 runtime required — plain pytest with numpy.
"""

import numpy as np
import pytest

from saltybot_segmentation.seg_utils import (
    TRAV_SIDEWALK, TRAV_GRASS, TRAV_ROAD, TRAV_OBSTACLE, TRAV_UNKNOWN,
    TRAV_TO_COST,
    cityscapes_to_traversability,
    traversability_to_costmap,
)


# ── cityscapes_to_traversability ──────────────────────────────────────────────

class TestCitiyscapesToTraversability:

    def test_sidewalk_maps_to_free(self):
        mask = np.array([[1]], dtype=np.uint8)  # Cityscapes: sidewalk
        result = cityscapes_to_traversability(mask)
        assert result[0, 0] == TRAV_SIDEWALK

    def test_road_maps_to_road(self):
        mask = np.array([[0]], dtype=np.uint8)  # Cityscapes: road
        result = cityscapes_to_traversability(mask)
        assert result[0, 0] == TRAV_ROAD

    def test_vegetation_maps_to_grass(self):
        mask = np.array([[8]], dtype=np.uint8)  # Cityscapes: vegetation
        result = cityscapes_to_traversability(mask)
        assert result[0, 0] == TRAV_GRASS

    def test_terrain_maps_to_grass(self):
        mask = np.array([[9]], dtype=np.uint8)  # Cityscapes: terrain
        result = cityscapes_to_traversability(mask)
        assert result[0, 0] == TRAV_GRASS

    def test_building_maps_to_obstacle(self):
        mask = np.array([[2]], dtype=np.uint8)  # Cityscapes: building
        result = cityscapes_to_traversability(mask)
        assert result[0, 0] == TRAV_OBSTACLE

    def test_person_maps_to_obstacle(self):
        mask = np.array([[11]], dtype=np.uint8)  # Cityscapes: person
        result = cityscapes_to_traversability(mask)
        assert result[0, 0] == TRAV_OBSTACLE

    def test_car_maps_to_obstacle(self):
        mask = np.array([[13]], dtype=np.uint8)  # Cityscapes: car
        result = cityscapes_to_traversability(mask)
        assert result[0, 0] == TRAV_OBSTACLE

    def test_sky_maps_to_unknown(self):
        mask = np.array([[10]], dtype=np.uint8)  # Cityscapes: sky
        result = cityscapes_to_traversability(mask)
        assert result[0, 0] == TRAV_UNKNOWN

    def test_all_dynamic_obstacles_are_lethal(self):
        # person=11, rider=12, car=13, truck=14, bus=15, train=16, motorcycle=17, bicycle=18
        for cid in [11, 12, 13, 14, 15, 16, 17, 18]:
            mask = np.array([[cid]], dtype=np.uint8)
            result = cityscapes_to_traversability(mask)
            assert result[0, 0] == TRAV_OBSTACLE, f"class {cid} should be OBSTACLE"

    def test_all_static_obstacles_are_lethal(self):
        # building=2, wall=3, fence=4, pole=5, traffic light=6, sign=7
        for cid in [2, 3, 4, 5, 6, 7]:
            mask = np.array([[cid]], dtype=np.uint8)
            result = cityscapes_to_traversability(mask)
            assert result[0, 0] == TRAV_OBSTACLE, f"class {cid} should be OBSTACLE"

    def test_out_of_range_id_clamped(self):
        """Class IDs beyond 18 should not crash — clamped to 18."""
        mask = np.array([[200]], dtype=np.uint8)
        result = cityscapes_to_traversability(mask)
        assert result[0, 0] in range(5)  # valid traversability class

    def test_multi_class_image(self):
        """A 2×3 image with mixed classes maps correctly."""
        mask = np.array([
            [1, 0, 8],   # sidewalk, road, vegetation
            [2, 11, 10], # building, person, sky
        ], dtype=np.uint8)
        result = cityscapes_to_traversability(mask)
        assert result[0, 0] == TRAV_SIDEWALK
        assert result[0, 1] == TRAV_ROAD
        assert result[0, 2] == TRAV_GRASS
        assert result[1, 0] == TRAV_OBSTACLE
        assert result[1, 1] == TRAV_OBSTACLE
        assert result[1, 2] == TRAV_UNKNOWN

    def test_output_dtype_is_uint8(self):
        mask = np.zeros((4, 4), dtype=np.uint8)
        result = cityscapes_to_traversability(mask)
        assert result.dtype == np.uint8

    def test_output_shape_preserved(self):
        mask = np.zeros((32, 64), dtype=np.uint8)
        result = cityscapes_to_traversability(mask)
        assert result.shape == (32, 64)


# ── traversability_to_costmap ─────────────────────────────────────────────────

class TestTraversabilityToCostmap:

    def test_sidewalk_is_free(self):
        mask = np.array([[TRAV_SIDEWALK]], dtype=np.uint8)
        cost = traversability_to_costmap(mask)
        assert cost[0, 0] == 0

    def test_road_has_high_cost(self):
        mask = np.array([[TRAV_ROAD]], dtype=np.uint8)
        cost = traversability_to_costmap(mask)
        assert 50 < cost[0, 0] < 100  # high but not lethal

    def test_grass_has_medium_cost(self):
        mask = np.array([[TRAV_GRASS]], dtype=np.uint8)
        cost = traversability_to_costmap(mask)
        grass_cost = TRAV_TO_COST[TRAV_GRASS]
        assert cost[0, 0] == grass_cost
        assert grass_cost < TRAV_TO_COST[TRAV_ROAD]  # grass < road cost

    def test_obstacle_is_lethal(self):
        mask = np.array([[TRAV_OBSTACLE]], dtype=np.uint8)
        cost = traversability_to_costmap(mask)
        assert cost[0, 0] == 100

    def test_unknown_is_neg1_by_default(self):
        mask = np.array([[TRAV_UNKNOWN]], dtype=np.uint8)
        cost = traversability_to_costmap(mask)
        assert cost[0, 0] == -1

    def test_unknown_as_obstacle_flag(self):
        mask = np.array([[TRAV_UNKNOWN]], dtype=np.uint8)
        cost = traversability_to_costmap(mask, unknown_as_obstacle=True)
        assert cost[0, 0] == 100

    def test_cost_ordering(self):
        """sidewalk < grass < road < obstacle."""
        assert TRAV_TO_COST[TRAV_SIDEWALK] < TRAV_TO_COST[TRAV_GRASS]
        assert TRAV_TO_COST[TRAV_GRASS]    < TRAV_TO_COST[TRAV_ROAD]
        assert TRAV_TO_COST[TRAV_ROAD]     < TRAV_TO_COST[TRAV_OBSTACLE]

    def test_output_dtype_is_int8(self):
        mask = np.zeros((4, 4), dtype=np.uint8)
        cost = traversability_to_costmap(mask)
        assert cost.dtype == np.int8

    def test_output_shape_preserved(self):
        mask = np.zeros((16, 32), dtype=np.uint8)
        cost = traversability_to_costmap(mask)
        assert cost.shape == (16, 32)

    def test_mixed_mask(self):
        mask = np.array([
            [TRAV_SIDEWALK, TRAV_ROAD],
            [TRAV_OBSTACLE, TRAV_UNKNOWN],
        ], dtype=np.uint8)
        cost = traversability_to_costmap(mask)
        assert cost[0, 0] == TRAV_TO_COST[TRAV_SIDEWALK]
        assert cost[0, 1] == TRAV_TO_COST[TRAV_ROAD]
        assert cost[1, 0] == TRAV_TO_COST[TRAV_OBSTACLE]
        assert cost[1, 1] == -1  # unknown
