"""
yolo_utils.py — YOLOv8 preprocessing / postprocessing helpers.

Shared by scene_detector_node and build_scene_trt script.
No ROS2 imports — unit-testable standalone.

YOLOv8n output format (TensorRT / ONNX):
  Shape: [1, 84, 8400]  (batch=1, 4_bbox + 80_classes, anchors)
  Coordinate order: cx, cy, w, h  (normalised 0–1 in 640×640 input space)

Custom-extended class set for scene understanding (appended after COCO 80):
  80 = door       (no COCO equivalent; detected via custom-trained head)
  81 = stairs     (custom)
  82 = wet_floor  (custom — yellow warning sign or wet-floor pattern)
The custom head is optional — if the engine only has 80 outputs the custom
classes simply won't appear.
"""

from __future__ import annotations
import numpy as np

# ── COCO 80 class names (index = COCO class id) ──────────────────────────────
COCO_NAMES = [
    'person','bicycle','car','motorcycle','airplane','bus','train','truck','boat',
    'traffic light','fire hydrant','stop sign','parking meter','bench','bird','cat',
    'dog','horse','sheep','cow','elephant','bear','zebra','giraffe','backpack',
    'umbrella','handbag','tie','suitcase','frisbee','skis','snowboard','sports ball',
    'kite','baseball bat','baseball glove','skateboard','surfboard','tennis racket',
    'bottle','wine glass','cup','fork','knife','spoon','bowl','banana','apple',
    'sandwich','orange','broccoli','carrot','hot dog','pizza','donut','cake',
    'chair','couch','potted plant','bed','dining table','toilet','tv','laptop',
    'mouse','remote','keyboard','cell phone','microwave','oven','toaster','sink',
    'refrigerator','book','clock','vase','scissors','teddy bear','hair drier',
    'toothbrush',
    # Custom extensions (indices 80+):
    'door', 'stairs', 'wet_floor',
]

# ── Scene-relevant class IDs ─────────────────────────────────────────────────
SCENE_CLASSES = {
    0:  'person',
    15: 'cat',
    16: 'dog',
    56: 'chair',
    57: 'couch',
    58: 'potted plant',
    60: 'dining table',
    61: 'toilet',
    62: 'tv',
    68: 'microwave',
    69: 'oven',
    71: 'sink',
    72: 'refrigerator',
    80: 'door',
    81: 'stairs',
    82: 'wet_floor',
}

# ── Hazard class mapping ──────────────────────────────────────────────────────
# SceneObject.HAZARD_* constants (must match msg definition)
HAZARD_NONE      = 0
HAZARD_STAIRS    = 1
HAZARD_DROP      = 2
HAZARD_WET_FLOOR = 3
HAZARD_GLASS_DOOR = 4
HAZARD_PET       = 5

HAZARD_BY_CLASS = {
    81: HAZARD_STAIRS,
    82: HAZARD_WET_FLOOR,
    15: HAZARD_PET,    # cat
    16: HAZARD_PET,    # dog
}

# ── Input preprocessing ──────────────────────────────────────────────────────
_YOLO_SIZE = 640


def letterbox(image: np.ndarray, size: int = _YOLO_SIZE, pad_value: int = 114):
    """
    Letterbox-resize BGR image to size×size.
    Returns (canvas, scale, pad_w, pad_h).
    """
    import cv2
    h, w = image.shape[:2]
    scale = min(size / w, size / h)
    new_w, new_h = int(round(w * scale)), int(round(h * scale))
    resized = cv2.resize(image, (new_w, new_h), interpolation=cv2.INTER_LINEAR)
    canvas = np.full((size, size, 3), pad_value, dtype=np.uint8)
    pad_w = (size - new_w) // 2
    pad_h = (size - new_h) // 2
    canvas[pad_h:pad_h + new_h, pad_w:pad_w + new_w] = resized
    return canvas, scale, pad_w, pad_h


def preprocess(bgr: np.ndarray) -> tuple[np.ndarray, float, int, int]:
    """
    Prepare a BGR frame for YOLOv8 inference.

    Returns:
        (blob, scale, pad_w, pad_h)
        blob: float32 CHW input, shape (1, 3, 640, 640), values 0–1
    """
    canvas, scale, pad_w, pad_h = letterbox(bgr)
    rgb = canvas[..., ::-1]                          # BGR → RGB
    blob = rgb.astype(np.float32) / 255.0
    blob = blob.transpose(2, 0, 1)[None]             # HWC → 1CHW
    blob = np.ascontiguousarray(blob)
    return blob, scale, pad_w, pad_h


def postprocess(
    raw: np.ndarray,
    scale: float,
    pad_w: int,
    pad_h: int,
    orig_w: int,
    orig_h: int,
    conf_thr: float = 0.35,
    iou_thr: float  = 0.45,
) -> list[dict]:
    """
    Parse YOLOv8 output tensor → list of detection dicts.

    Args:
        raw:    shape [1, num_classes+4, 8400]  (as returned by TRT/ONNX)
        scale, pad_w, pad_h:  from preprocess()
        orig_w, orig_h:       original frame dimensions

    Returns:
        List of dicts with keys: class_id, class_name, conf, x1, y1, x2, y2
    """
    preds = raw[0]          # (84+, 8400)
    # YOLOv8 output layout: rows 0-3 = cx,cy,w,h; rows 4+ = class scores
    boxes_xywh  = preds[:4].T     # (8400, 4)
    class_probs = preds[4:].T     # (8400, N_classes)

    class_ids   = class_probs.argmax(axis=1)
    confidences = class_probs[np.arange(len(class_probs)), class_ids]

    # Filter by confidence
    mask = confidences > conf_thr
    if not mask.any():
        return []

    boxes_xywh  = boxes_xywh[mask]
    confidences = confidences[mask]
    class_ids   = class_ids[mask]

    # cx,cy,w,h → x1,y1,x2,y2 (letterbox space)
    cx, cy, w, h = boxes_xywh.T
    x1 = cx - w / 2
    y1 = cy - h / 2
    x2 = cx + w / 2
    y2 = cy + h / 2
    boxes_xyxy = np.stack([x1, y1, x2, y2], axis=1)

    # NMS
    kept = _nms(boxes_xyxy, confidences, iou_thr)

    results = []
    for i in kept:
        cx1, cy1, cx2, cy2 = _remap(
            boxes_xyxy[i], scale, pad_w, pad_h, orig_w, orig_h
        )
        cid = int(class_ids[i])
        results.append({
            'class_id':   cid,
            'class_name': COCO_NAMES[cid] if cid < len(COCO_NAMES) else str(cid),
            'conf':       float(confidences[i]),
            'x1': cx1, 'y1': cy1, 'x2': cx2, 'y2': cy2,
        })
    return results


def _nms(boxes: np.ndarray, scores: np.ndarray, iou_thr: float) -> list[int]:
    x1, y1, x2, y2 = boxes[:, 0], boxes[:, 1], boxes[:, 2], boxes[:, 3]
    areas = np.maximum(0, x2 - x1) * np.maximum(0, y2 - y1)
    order = scores.argsort()[::-1]
    keep  = []
    while len(order):
        i = order[0]
        keep.append(int(i))
        if len(order) == 1:
            break
        ix1 = np.maximum(x1[i], x1[order[1:]])
        iy1 = np.maximum(y1[i], y1[order[1:]])
        ix2 = np.minimum(x2[i], x2[order[1:]])
        iy2 = np.minimum(y2[i], y2[order[1:]])
        inter = np.maximum(0, ix2 - ix1) * np.maximum(0, iy2 - iy1)
        iou   = inter / (areas[i] + areas[order[1:]] - inter + 1e-6)
        order = order[1:][iou < iou_thr]
    return keep


def _remap(box, scale, pad_w, pad_h, orig_w, orig_h):
    x1, y1, x2, y2 = box
    x1 = float(np.clip((x1 - pad_w) / scale, 0, orig_w))
    y1 = float(np.clip((y1 - pad_h) / scale, 0, orig_h))
    x2 = float(np.clip((x2 - pad_w) / scale, 0, orig_w))
    y2 = float(np.clip((y2 - pad_h) / scale, 0, orig_h))
    return x1, y1, x2, y2
