"""
seg_utils.py — Pure helpers for semantic segmentation and traversability mapping.

No ROS2, TensorRT, or PyTorch imports — fully testable without GPU or ROS2 install.

Cityscapes 19-class → 5-class traversability mapping
──────────────────────────────────────────────────────
  Cityscapes ID → Traversability class → OccupancyGrid value

  0  road         → ROAD      → 90   (high cost — don't drive on road)
  1  sidewalk     → SIDEWALK  → 0    (free — preferred surface)
  2  building     → OBSTACLE  → 100  (lethal)
  3  wall         → OBSTACLE  → 100
  4  fence        → OBSTACLE  → 100
  5  pole         → OBSTACLE  → 100
  6  traffic light→ OBSTACLE  → 100
  7  traffic sign → OBSTACLE  → 100
  8  vegetation   → GRASS     → 50   (medium — can traverse slowly)
  9  terrain      → GRASS     → 50
  10 sky          → UNKNOWN   → -1   (unknown)
  11 person       → OBSTACLE  → 100  (lethal — safety critical)
  12 rider        → OBSTACLE  → 100
  13 car          → OBSTACLE  → 100
  14 truck        → OBSTACLE  → 100
  15 bus          → OBSTACLE  → 100
  16 train        → OBSTACLE  → 100
  17 motorcycle   → OBSTACLE  → 100
  18 bicycle      → OBSTACLE  → 100

Segmentation mask class IDs (published on /segmentation/mask)
─────────────────────────────────────────────────────────────
  TRAV_SIDEWALK = 0
  TRAV_GRASS    = 1
  TRAV_ROAD     = 2
  TRAV_OBSTACLE = 3
  TRAV_UNKNOWN  = 4
"""

import numpy as np

# ── Traversability class constants ────────────────────────────────────────────

TRAV_SIDEWALK = 0
TRAV_GRASS    = 1
TRAV_ROAD     = 2
TRAV_OBSTACLE = 3
TRAV_UNKNOWN  = 4

NUM_TRAV_CLASSES = 5

# ── OccupancyGrid cost values for each traversability class ───────────────────
# Nav2 OccupancyGrid: 0=free, 1-99=cost, 100=lethal, -1=unknown

TRAV_TO_COST = {
    TRAV_SIDEWALK: 0,     # free — preferred surface
    TRAV_GRASS:    50,    # medium cost — traversable but non-preferred
    TRAV_ROAD:     90,    # high cost — avoid but can cross
    TRAV_OBSTACLE: 100,   # lethal — never enter
    TRAV_UNKNOWN:  -1,    # unknown — Nav2 treats as unknown cell
}

# ── Cityscapes 19-class → traversability lookup table ─────────────────────────
# Index = Cityscapes training ID (0–18)

_CITYSCAPES_TO_TRAV = np.array([
    TRAV_ROAD,      # 0  road
    TRAV_SIDEWALK,  # 1  sidewalk
    TRAV_OBSTACLE,  # 2  building
    TRAV_OBSTACLE,  # 3  wall
    TRAV_OBSTACLE,  # 4  fence
    TRAV_OBSTACLE,  # 5  pole
    TRAV_OBSTACLE,  # 6  traffic light
    TRAV_OBSTACLE,  # 7  traffic sign
    TRAV_GRASS,     # 8  vegetation
    TRAV_GRASS,     # 9  terrain
    TRAV_UNKNOWN,   # 10 sky
    TRAV_OBSTACLE,  # 11 person
    TRAV_OBSTACLE,  # 12 rider
    TRAV_OBSTACLE,  # 13 car
    TRAV_OBSTACLE,  # 14 truck
    TRAV_OBSTACLE,  # 15 bus
    TRAV_OBSTACLE,  # 16 train
    TRAV_OBSTACLE,  # 17 motorcycle
    TRAV_OBSTACLE,  # 18 bicycle
], dtype=np.uint8)

# ── Visualisation colour map (BGR, for cv2.imshow / debug image) ──────────────
# One colour per traversability class

TRAV_COLORMAP_BGR = np.array([
    [128, 64, 128],   # SIDEWALK  — purple
    [152, 251, 152],  # GRASS     — pale green
    [128,   0, 128],  # ROAD      — dark magenta
    [ 60,  20, 220],  # OBSTACLE  — red-ish
    [  0,   0,   0],  # UNKNOWN   — black
], dtype=np.uint8)


# ── Core mapping functions ────────────────────────────────────────────────────

def cityscapes_to_traversability(mask: np.ndarray) -> np.ndarray:
    """
    Map a Cityscapes 19-class segmentation mask to traversability classes.

    Parameters
    ----------
    mask : np.ndarray, shape (H, W), dtype uint8
        Cityscapes class IDs in range [0, 18]. Values outside range are
        mapped to TRAV_UNKNOWN.

    Returns
    -------
    trav_mask : np.ndarray, shape (H, W), dtype uint8
        Traversability class IDs in range [0, 4].
    """
    # Clamp out-of-range IDs to a safe default (unknown)
    clipped = np.clip(mask.astype(np.int32), 0, len(_CITYSCAPES_TO_TRAV) - 1)
    return _CITYSCAPES_TO_TRAV[clipped]


def traversability_to_costmap(
    trav_mask: np.ndarray,
    unknown_as_obstacle: bool = False,
) -> np.ndarray:
    """
    Convert a traversability mask to Nav2 OccupancyGrid int8 cost values.

    Parameters
    ----------
    trav_mask : np.ndarray, shape (H, W), dtype uint8
        Traversability class IDs (TRAV_* constants).
    unknown_as_obstacle : bool
        If True, TRAV_UNKNOWN cells are set to 100 (lethal) instead of -1.
        Useful in dense urban environments where unknowns are likely walls.

    Returns
    -------
    cost_map : np.ndarray, shape (H, W), dtype int8
        OccupancyGrid cost values: 0=free, 1-99=cost, 100=lethal, -1=unknown.
    """
    H, W = trav_mask.shape
    cost_map = np.full((H, W), -1, dtype=np.int8)

    for trav_class, cost in TRAV_TO_COST.items():
        if trav_class == TRAV_UNKNOWN and unknown_as_obstacle:
            cost = 100
        cost_map[trav_mask == trav_class] = cost

    return cost_map


def letterbox(
    image: np.ndarray,
    target_w: int,
    target_h: int,
    pad_value: int = 114,
) -> tuple:
    """
    Letterbox-resize an image to (target_w, target_h) preserving aspect ratio.

    Parameters
    ----------
    image : np.ndarray, shape (H, W, C), dtype uint8
    target_w, target_h : int
        Target width and height.
    pad_value : int
        Padding fill value.

    Returns
    -------
    (canvas, scale, pad_left, pad_top) : tuple
        canvas   : np.ndarray (target_h, target_w, C) — resized + padded image
        scale    : float — scale factor applied (min of w_scale, h_scale)
        pad_left : int — horizontal padding (pixels from left)
        pad_top  : int — vertical padding (pixels from top)
    """
    import cv2  # lazy import — keeps module importable without cv2 in tests

    src_h, src_w = image.shape[:2]
    scale = min(target_w / src_w, target_h / src_h)
    new_w = int(round(src_w * scale))
    new_h = int(round(src_h * scale))

    resized = cv2.resize(image, (new_w, new_h), interpolation=cv2.INTER_LINEAR)

    pad_left = (target_w - new_w) // 2
    pad_top  = (target_h - new_h) // 2

    canvas = np.full((target_h, target_w, image.shape[2]), pad_value, dtype=np.uint8)
    canvas[pad_top:pad_top + new_h, pad_left:pad_left + new_w] = resized

    return canvas, scale, pad_left, pad_top


def unpad_mask(
    mask: np.ndarray,
    orig_h: int,
    orig_w: int,
    scale: float,
    pad_left: int,
    pad_top: int,
) -> np.ndarray:
    """
    Remove letterbox padding from a segmentation mask and resize to original.

    Parameters
    ----------
    mask : np.ndarray, shape (target_h, target_w), dtype uint8
        Segmentation output from model at letterboxed resolution.
    orig_h, orig_w : int
        Original image dimensions (before letterboxing).
    scale, pad_left, pad_top : from letterbox()

    Returns
    -------
    restored : np.ndarray, shape (orig_h, orig_w), dtype uint8
    """
    import cv2

    new_h = int(round(orig_h * scale))
    new_w = int(round(orig_w * scale))

    cropped = mask[pad_top:pad_top + new_h, pad_left:pad_left + new_w]
    restored = cv2.resize(cropped, (orig_w, orig_h), interpolation=cv2.INTER_NEAREST)
    return restored


def preprocess_for_inference(
    image_bgr: np.ndarray,
    input_w: int = 512,
    input_h: int = 256,
) -> tuple:
    """
    Preprocess a BGR image for BiSeNetV2 / DDRNet inference.

    Applies letterboxing, BGR→RGB conversion, ImageNet normalisation,
    HWC→NCHW layout, and float32 conversion.

    Returns
    -------
    (blob, scale, pad_left, pad_top) : tuple
        blob : np.ndarray (1, 3, input_h, input_w) float32 — model input
        scale, pad_left, pad_top : for unpad_mask()
    """
    import cv2

    canvas, scale, pad_left, pad_top = letterbox(image_bgr, input_w, input_h)

    # BGR → RGB
    rgb = cv2.cvtColor(canvas, cv2.COLOR_BGR2RGB)

    # ImageNet normalisation
    mean = np.array([0.485, 0.456, 0.406], dtype=np.float32)
    std  = np.array([0.229, 0.224, 0.225], dtype=np.float32)
    img  = rgb.astype(np.float32) / 255.0
    img  = (img - mean) / std

    # HWC → NCHW
    blob = np.ascontiguousarray(img.transpose(2, 0, 1)[np.newaxis])

    return blob, scale, pad_left, pad_top


def colorise_traversability(trav_mask: np.ndarray) -> np.ndarray:
    """
    Convert traversability mask to an RGB visualisation image.

    Returns
    -------
    vis : np.ndarray (H, W, 3) uint8 BGR — suitable for cv2.imshow / ROS Image
    """
    clipped = np.clip(trav_mask.astype(np.int32), 0, NUM_TRAV_CLASSES - 1)
    return TRAV_COLORMAP_BGR[clipped]
