"""
scene_detector_node.py — Semantic scene understanding: YOLOv8-nano TRT FP16.

Pipeline per frame:
  1. Receive synchronized RGB + depth from RealSense D435i.
  2. YOLOv8-nano TRT FP16 inference (target ≥15 FPS).
  3. Filter scene-relevant classes (chairs, tables, doors, stairs, pets, …).
  4. Back-project bounding box centres to 3D via aligned depth.
  5. Classify hazard type per object (stairs, drop, wet floor, glass door, pet).
  6. Run depth-image hazard scan (HazardClassifier) for drop/wet not tied to a bbox.
  7. Classify room type (rule_based or MobileNetV2 TRT).
  8. Publish results.

Subscribes:
  /camera/color/image_raw            sensor_msgs/Image       30 Hz RGB
  /camera/depth/image_rect_raw       sensor_msgs/Image       30 Hz float32 depth (m)
  /camera/color/camera_info          sensor_msgs/CameraInfo  camera intrinsics

Publishes:
  /social/scene/objects              SceneObjectArray        per-frame detections
  /social/scene/room_type            RoomClassification      ~2 Hz (decimated)
  /social/scene/hazards              SceneObjectArray        hazard subset only
  /social/scene/debug_image          sensor_msgs/Image       annotated RGB (lazy)

Parameters:
  engine_path       str  '/config/yolov8n_scene.engine'
  onnx_path         str  '/config/yolov8n_scene.onnx'
  room_engine_path  str  '/config/room_classifier.engine'
  room_onnx_path    str  '/config/room_classifier.onnx'
  conf_threshold    float 0.35
  iou_threshold     float 0.45
  publish_debug     bool  false
  room_update_hz    float 2.0   (room classification rate, decimated from detection rate)
  max_depth_m       float 6.0
"""

import time
import math

import rclpy
from rclpy.node import Node
from rclpy.qos import QoSProfile, ReliabilityPolicy, HistoryPolicy, DurabilityPolicy
import message_filters
import cv2
import numpy as np
from cv_bridge import CvBridge

from sensor_msgs.msg import Image, CameraInfo
from geometry_msgs.msg import PointStamped, Point
from vision_msgs.msg import BoundingBox2D, Pose2D

import tf2_ros
import tf2_geometry_msgs  # noqa: F401

from saltybot_scene_msgs.msg import (
    SceneObject, SceneObjectArray, RoomClassification
)

from .yolo_utils import (
    preprocess, postprocess, SCENE_CLASSES, HAZARD_BY_CLASS,
    HAZARD_NONE, HAZARD_STAIRS, HAZARD_PET
)
from .room_classifier import RoomClassifier, ROOM_NAMES
from .hazard_classifier import HazardClassifier


_LATCHED_QOS = QoSProfile(
    reliability=ReliabilityPolicy.RELIABLE,
    history=HistoryPolicy.KEEP_LAST,
    depth=1,
    durability=DurabilityPolicy.TRANSIENT_LOCAL,
)
_SENSOR_QOS = QoSProfile(
    reliability=ReliabilityPolicy.BEST_EFFORT,
    history=HistoryPolicy.KEEP_LAST,
    depth=5,
)


# ── TensorRT backend (identical pattern to person_detector_node) ──────────────

class _TRTBackend:
    def __init__(self, engine_path: str):
        import tensorrt as trt
        import pycuda.driver as cuda
        import pycuda.autoinit  # noqa: F401
        self._cuda = cuda
        logger = trt.Logger(trt.Logger.WARNING)
        with open(engine_path, 'rb') as f, trt.Runtime(logger) as rt:
            self._engine = rt.deserialize_cuda_engine(f.read())
        self._ctx = self._engine.create_execution_context()
        self._inputs, self._outputs, self._bindings = [], [], []
        for i in range(self._engine.num_io_tensors):
            name  = self._engine.get_tensor_name(i)
            dtype = trt.nptype(self._engine.get_tensor_dtype(name))
            shape = tuple(self._engine.get_tensor_shape(name))
            nbytes = int(np.prod(shape)) * np.dtype(dtype).itemsize
            host   = cuda.pagelocked_empty(shape, dtype)
            device = cuda.mem_alloc(nbytes)
            self._bindings.append(int(device))
            if self._engine.get_tensor_mode(name) == trt.TensorIOMode.INPUT:
                self._inputs.append({'host': host, 'device': device})
            else:
                self._outputs.append({'host': host, 'device': device})
        self._stream = cuda.Stream()

    def infer(self, blob: np.ndarray) -> list[np.ndarray]:
        np.copyto(self._inputs[0]['host'], blob.ravel())
        self._cuda.memcpy_htod_async(
            self._inputs[0]['device'], self._inputs[0]['host'], self._stream
        )
        self._ctx.execute_async_v2(self._bindings, self._stream.handle)
        for out in self._outputs:
            self._cuda.memcpy_dtoh_async(out['host'], out['device'], self._stream)
        self._stream.synchronize()
        return [out['host'].copy() for out in self._outputs]


class _ONNXBackend:
    def __init__(self, onnx_path: str):
        import onnxruntime as ort
        self._session   = ort.InferenceSession(
            onnx_path, providers=['CUDAExecutionProvider', 'CPUExecutionProvider']
        )
        self._input_name = self._session.get_inputs()[0].name

    def infer(self, blob: np.ndarray) -> list[np.ndarray]:
        return self._session.run(None, {self._input_name: blob})


# ── Node ──────────────────────────────────────────────────────────────────────

class SceneDetectorNode(Node):

    def __init__(self):
        super().__init__('scene_detector')

        # ── Parameters ──────────────────────────────────────────────────────
        self.declare_parameter('engine_path',      '/config/yolov8n_scene.engine')
        self.declare_parameter('onnx_path',        '/config/yolov8n_scene.onnx')
        self.declare_parameter('room_engine_path', '/config/room_classifier.engine')
        self.declare_parameter('room_onnx_path',   '/config/room_classifier.onnx')
        self.declare_parameter('conf_threshold',   0.35)
        self.declare_parameter('iou_threshold',    0.45)
        self.declare_parameter('publish_debug',    False)
        self.declare_parameter('room_update_hz',   2.0)
        self.declare_parameter('max_depth_m',      6.0)

        engine_path  = self.get_parameter('engine_path').value
        onnx_path    = self.get_parameter('onnx_path').value
        self._conf   = self.get_parameter('conf_threshold').value
        self._iou    = self.get_parameter('iou_threshold').value
        self._debug  = self.get_parameter('publish_debug').value
        room_hz      = self.get_parameter('room_update_hz').value
        self._max_d  = self.get_parameter('max_depth_m').value

        # ── Inference backend ────────────────────────────────────────────────
        self._backend = None
        for cls, path in [(_TRTBackend, engine_path), (_ONNXBackend, onnx_path)]:
            if path:
                try:
                    self._backend = cls(path)
                    self.get_logger().info(
                        f'Loaded {cls.__name__} from {path}'
                    )
                    break
                except Exception as exc:
                    self.get_logger().warning(f'{cls.__name__} load failed: {exc}')

        if self._backend is None:
            self.get_logger().error(
                'No inference backend loaded — set engine_path or onnx_path param. '
                'Publishing empty detections.'
            )

        # ── Room classifier ──────────────────────────────────────────────────
        room_engine = self.get_parameter('room_engine_path').value
        room_onnx   = self.get_parameter('room_onnx_path').value
        self._room_clf = RoomClassifier(
            engine_path=room_engine, onnx_path=room_onnx
        )

        # ── Hazard classifier ────────────────────────────────────────────────
        self._hazard_clf = HazardClassifier()

        # ── State ────────────────────────────────────────────────────────────
        self._bridge         = CvBridge()
        self._camera_k: np.ndarray | None = None
        self._frame_count    = 0
        self._room_period    = max(1, int(round(30.0 / max(room_hz, 0.1))))
        self._last_room: RoomClassification | None = None

        # ── TF ───────────────────────────────────────────────────────────────
        self._tf_buffer    = tf2_ros.Buffer()
        self._tf_listener  = tf2_ros.TransformListener(self._tf_buffer, self)

        # ── Synchronized subscriptions ────────────────────────────────────────
        rgb_sub   = message_filters.Subscriber(
            self, Image, '/camera/color/image_raw',       qos_profile=_SENSOR_QOS
        )
        depth_sub = message_filters.Subscriber(
            self, Image, '/camera/depth/image_rect_raw',  qos_profile=_SENSOR_QOS
        )
        self._sync = message_filters.ApproximateTimeSynchronizer(
            [rgb_sub, depth_sub], queue_size=5, slop=0.05
        )
        self._sync.registerCallback(self._on_rgbd)

        self.create_subscription(
            CameraInfo, '/camera/color/camera_info', self._on_camera_info, _LATCHED_QOS
        )

        # ── Publishers ───────────────────────────────────────────────────────
        self._obj_pub     = self.create_publisher(SceneObjectArray, '/social/scene/objects',   10)
        self._hazard_pub  = self.create_publisher(SceneObjectArray, '/social/scene/hazards',   10)
        self._room_pub    = self.create_publisher(RoomClassification, '/social/scene/room_type', 10)
        if self._debug:
            self._debug_pub = self.create_publisher(Image, '/social/scene/debug_image', 5)

        self.get_logger().info(
            f'scene_detector ready — backend={"TRT" if isinstance(self._backend, _TRTBackend) else "ONNX" if self._backend else "none"}'
            f', room_clf={self._room_clf._method}'
        )

    # ── Callbacks ────────────────────────────────────────────────────────────

    def _on_camera_info(self, msg: CameraInfo) -> None:
        self._camera_k = np.array(msg.k, dtype=np.float64)

    def _on_rgbd(self, rgb_msg: Image, depth_msg: Image) -> None:
        t0 = time.monotonic()
        self._frame_count += 1

        # Decode images
        try:
            bgr   = self._bridge.imgmsg_to_cv2(rgb_msg,   desired_encoding='bgr8')
            depth = self._bridge.imgmsg_to_cv2(depth_msg, desired_encoding='32FC1')
        except Exception as exc:
            self.get_logger().warning(f'Image decode error: {exc}')
            return

        orig_h, orig_w = bgr.shape[:2]

        # ── YOLOv8 inference ─────────────────────────────────────────────────
        detections = []
        infer_ms   = 0.0
        if self._backend is not None:
            blob, scale, pad_w, pad_h = preprocess(bgr)
            t_infer = time.monotonic()
            raw     = self._backend.infer(blob)
            infer_ms = (time.monotonic() - t_infer) * 1000.0
            detections = postprocess(
                raw[0], scale, pad_w, pad_h, orig_w, orig_h,
                conf_thr=self._conf, iou_thr=self._iou,
            )
            # Filter to scene-relevant classes only
            detections = [d for d in detections if d['class_id'] in SCENE_CLASSES]

        # ── Depth-based hazard scan (independent of YOLO) ────────────────────
        depth_hazards = self._hazard_clf.analyse(depth, bgr)

        # ── Build SceneObject messages ────────────────────────────────────────
        objects = []
        for det in detections:
            obj = self._build_scene_object(det, depth, rgb_msg.header)
            objects.append(obj)

        # Inject depth-detected hazards as pseudo-objects (no bbox)
        for hazard_type, conf, desc in depth_hazards:
            # Check we don't already have a YOLO object for the same hazard type
            already = any(o.hazard_type == hazard_type for o in objects)
            if not already:
                pseudo = SceneObject()
                pseudo.class_id   = 255          # sentinel: depth-only hazard
                pseudo.class_name = desc
                pseudo.confidence = conf
                pseudo.hazard_type = hazard_type
                objects.append(pseudo)

        # ── Publish objects ───────────────────────────────────────────────────
        arr = SceneObjectArray()
        arr.header       = rgb_msg.header
        arr.objects      = objects
        arr.frame_count  = self._frame_count
        arr.inference_ms = infer_ms
        self._obj_pub.publish(arr)

        # ── Publish hazard subset ────────────────────────────────────────────
        hazard_objs = [o for o in objects if o.hazard_type != 0]
        if hazard_objs:
            harr = SceneObjectArray()
            harr.header      = rgb_msg.header
            harr.objects     = hazard_objs
            harr.frame_count = self._frame_count
            self._hazard_pub.publish(harr)

        # ── Room classification (decimated) ───────────────────────────────────
        if self._frame_count % self._room_period == 0:
            class_ids   = [d['class_id']   for d in detections]
            confs       = [d['conf']        for d in detections]
            room_type, room_name, confidence, scores, method = \
                self._room_clf.classify(bgr, class_ids, confs)
            room_msg = RoomClassification()
            room_msg.header       = rgb_msg.header
            room_msg.room_type    = room_type
            room_msg.room_name    = room_name
            room_msg.confidence   = confidence
            room_msg.class_scores = scores
            room_msg.method       = method
            self._room_pub.publish(room_msg)
            self._last_room = room_msg

        # ── Debug image ───────────────────────────────────────────────────────
        if self._debug and hasattr(self, '_debug_pub'):
            self._debug_pub.publish(
                self._bridge.cv2_to_imgmsg(
                    self._draw_debug(bgr, objects), encoding='bgr8'
                )
            )

        total_ms = (time.monotonic() - t0) * 1000.0
        if self._frame_count % 30 == 0:
            fps = 1000.0 / max(total_ms, 1)
            room = self._last_room.room_name if self._last_room else 'unknown'
            self.get_logger().info(
                f'[scene] {fps:.1f} FPS  infer={infer_ms:.1f}ms  '
                f'{len(objects)} objs  room={room}'
            )

    # ── Helpers ───────────────────────────────────────────────────────────────

    def _build_scene_object(
        self, det: dict, depth: np.ndarray, header
    ) -> SceneObject:
        obj = SceneObject()
        obj.class_id   = det['class_id']
        obj.class_name = det['class_name']
        obj.confidence = det['conf']

        # Bounding box
        cx = (det['x1'] + det['x2']) / 2
        cy = (det['y1'] + det['y2']) / 2
        bw = det['x2'] - det['x1']
        bh = det['y2'] - det['y1']
        obj.bbox.center.position.x = cx
        obj.bbox.center.position.y = cy
        obj.bbox.size_x = bw
        obj.bbox.size_y = bh

        # 3-D position
        if self._camera_k is not None:
            h, w = depth.shape
            cx_i, cy_i = int(np.clip(cx, 0, w - 1)), int(np.clip(cy, 0, h - 1))
            d = self._sample_depth(depth, cx_i, cy_i)
            if d > 0.1:
                obj.distance_m = d
                K = self._camera_k
                X = (cx - K[2]) * d / K[0]
                Y = (cy - K[5]) * d / K[4]
                pt_cam = PointStamped()
                pt_cam.header         = header
                pt_cam.point.x        = X
                pt_cam.point.y        = Y
                pt_cam.point.z        = d
                obj.position_3d       = pt_cam
                # Optionally transform to base_link
                try:
                    pt_bl = self._tf_buffer.transform(pt_cam, 'base_link',
                                                       rclpy.duration.Duration(seconds=0.05))
                    obj.position_3d = pt_bl
                except Exception:
                    pass   # keep camera-frame position

        # Hazard classification
        obj.hazard_type = HAZARD_BY_CLASS.get(det['class_id'], HAZARD_NONE)

        return obj

    @staticmethod
    def _sample_depth(depth: np.ndarray, u: int, v: int, window: int = 9) -> float:
        h, w = depth.shape
        half = window // 2
        patch = depth[
            max(0, v - half): min(h, v + half + 1),
            max(0, u - half): min(w, u + half + 1),
        ]
        valid = patch[(patch > 0.1) & (patch < 8.0)]
        return float(np.median(valid)) if len(valid) > 0 else 0.0

    def _draw_debug(self, bgr: np.ndarray, objects: list) -> np.ndarray:
        img = bgr.copy()
        colours = {
            0: (0, 200, 0), 1: (0, 0, 255), 2: (0, 0, 255),
            3: (0, 128, 255), 4: (0, 0, 255), 5: (255, 128, 0),
        }
        for obj in objects:
            if obj.bbox.size_x < 1:
                continue
            cx = int(obj.bbox.center.position.x)
            cy = int(obj.bbox.center.position.y)
            x1 = int(cx - obj.bbox.size_x / 2)
            y1 = int(cy - obj.bbox.size_y / 2)
            x2 = int(cx + obj.bbox.size_x / 2)
            y2 = int(cy + obj.bbox.size_y / 2)
            colour = colours.get(obj.hazard_type, (0, 200, 0))
            cv2.rectangle(img, (x1, y1), (x2, y2), colour, 2)
            label = f'{obj.class_name} {obj.confidence:.0%}'
            if obj.distance_m > 0:
                label += f' {obj.distance_m:.1f}m'
            cv2.putText(img, label, (x1, y1 - 5),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, colour, 1)
        return img


def main(args=None):
    rclpy.init(args=args)
    node = SceneDetectorNode()
    try:
        rclpy.spin(node)
    finally:
        node.destroy_node()
        rclpy.shutdown()


if __name__ == '__main__':
    main()
