"""
sidewalk_seg_node.py — Semantic sidewalk segmentation node for SaltyBot.

Subscribes:
  /camera/color/image_raw   (sensor_msgs/Image) — RealSense front RGB

Publishes:
  /segmentation/mask        (sensor_msgs/Image, mono8) — traversability class per pixel
  /segmentation/costmap     (nav_msgs/OccupancyGrid)   — Nav2 traversability scores
  /segmentation/debug_image (sensor_msgs/Image, bgr8)  — colour-coded visualisation

Model backend (auto-selected, priority order)
─────────────────────────────────────────────
  1. TensorRT FP16 engine (.engine file)   — fastest, Orin GPU
  2. ONNX Runtime (CUDA provider)          — fallback, still GPU
  3. ONNX Runtime (CPU provider)           — last resort

Build TRT engine:
  python3 /opt/ros/humble/share/saltybot_segmentation/scripts/build_engine.py

Supported model architectures
──────────────────────────────
  BiSeNetV2   — Cityscapes 72.6 mIoU, ~50fps @ 512×256 on Orin
  DDRNet-23   — Cityscapes 79.5 mIoU, ~40fps @ 512×256 on Orin

Performance target: >15fps at 512×256 on Jetson Orin Nano Super.

Input:  512×256 RGB (letterboxed from 640×480 RealSense color)
Output: 512×256 per-pixel traversability class ID (uint8)
"""

import threading
import time

import cv2
import numpy as np
import rclpy
from cv_bridge import CvBridge
from nav_msgs.msg import OccupancyGrid, MapMetaData
from rclpy.node import Node
from rclpy.qos import QoSProfile, ReliabilityPolicy, HistoryPolicy, DurabilityPolicy
from sensor_msgs.msg import Image
from std_msgs.msg import Header

from saltybot_segmentation.seg_utils import (
    cityscapes_to_traversability,
    traversability_to_costmap,
    preprocess_for_inference,
    unpad_mask,
    colorise_traversability,
)

# ── TensorRT backend ──────────────────────────────────────────────────────────

try:
    import tensorrt as trt
    import pycuda.driver as cuda
    import pycuda.autoinit  # noqa: F401 — initialises CUDA context
    _TRT_AVAILABLE = True
except ImportError:
    _TRT_AVAILABLE = False


class _TRTBackend:
    """TensorRT FP16 inference backend."""

    def __init__(self, engine_path: str, logger):
        self._logger = logger
        trt_logger = trt.Logger(trt.Logger.WARNING)
        with open(engine_path, "rb") as f:
            serialized = f.read()
        runtime = trt.Runtime(trt_logger)
        self._engine  = runtime.deserialize_cuda_engine(serialized)
        self._context = self._engine.create_execution_context()
        self._stream  = cuda.Stream()

        # Allocate host + device buffers
        self._bindings = []
        self._host_in  = None
        self._dev_in   = None
        self._host_out = None
        self._dev_out  = None

        for i in range(self._engine.num_bindings):
            shape = self._engine.get_binding_shape(i)
            size  = int(np.prod(shape))
            dtype = trt.nptype(self._engine.get_binding_dtype(i))
            host_buf = cuda.pagelocked_empty(size, dtype)
            dev_buf  = cuda.mem_alloc(host_buf.nbytes)
            self._bindings.append(int(dev_buf))
            if self._engine.binding_is_input(i):
                self._host_in, self._dev_in = host_buf, dev_buf
                self._in_shape = shape
            else:
                self._host_out, self._dev_out = host_buf, dev_buf
                self._out_shape = shape

        logger.info(f"TRT engine loaded: in={self._in_shape} out={self._out_shape}")

    def infer(self, blob: np.ndarray) -> np.ndarray:
        np.copyto(self._host_in, blob.ravel())
        cuda.memcpy_htod_async(self._dev_in, self._host_in, self._stream)
        self._context.execute_async_v2(self._bindings, self._stream.handle)
        cuda.memcpy_dtoh_async(self._host_out, self._dev_out, self._stream)
        self._stream.synchronize()
        return self._host_out.reshape(self._out_shape)


# ── ONNX Runtime backend ──────────────────────────────────────────────────────

try:
    import onnxruntime as ort
    _ONNX_AVAILABLE = True
except ImportError:
    _ONNX_AVAILABLE = False


class _ONNXBackend:
    """ONNX Runtime inference backend (CUDA or CPU)."""

    def __init__(self, onnx_path: str, logger):
        providers = []
        if _ONNX_AVAILABLE:
            available = ort.get_available_providers()
            if "CUDAExecutionProvider" in available:
                providers.append("CUDAExecutionProvider")
                logger.info("ONNX Runtime: using CUDA provider")
            else:
                logger.warn("ONNX Runtime: CUDA not available, falling back to CPU")
            providers.append("CPUExecutionProvider")

        self._session   = ort.InferenceSession(onnx_path, providers=providers)
        self._input_name  = self._session.get_inputs()[0].name
        self._output_name = self._session.get_outputs()[0].name
        logger.info(f"ONNX model loaded: {onnx_path}")

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


# ── ROS2 Node ─────────────────────────────────────────────────────────────────

class SidewalkSegNode(Node):

    def __init__(self):
        super().__init__("sidewalk_seg")

        # ── Parameters ────────────────────────────────────────────────────────
        self.declare_parameter("engine_path",
            "/mnt/nvme/saltybot/models/bisenetv2_cityscapes_512x256.engine")
        self.declare_parameter("onnx_path",
            "/mnt/nvme/saltybot/models/bisenetv2_cityscapes_512x256.onnx")
        self.declare_parameter("input_width",        512)
        self.declare_parameter("input_height",       256)
        self.declare_parameter("process_every_n",    2)     # skip frames to save GPU
        self.declare_parameter("publish_debug_image", False)
        self.declare_parameter("unknown_as_obstacle", False)
        self.declare_parameter("costmap_resolution",  0.05)  # metres/cell
        self.declare_parameter("costmap_range_m",     5.0)   # forward projection range
        self.declare_parameter("camera_height_m",     0.15)  # RealSense mount height
        self.declare_parameter("camera_pitch_deg",    0.0)   # tilt angle

        self._p = self._load_params()

        # ── Backend selection (TRT preferred) ──────────────────────────────────
        self._backend = None
        self._init_backend()

        # ── Runtime state ─────────────────────────────────────────────────────
        self._bridge        = CvBridge()
        self._frame_count   = 0
        self._last_mask     = None
        self._last_mask_lock = threading.Lock()
        self._t_infer        = 0.0

        # ── QoS: sensor data (best-effort, keep-last=1) ────────────────────────
        sensor_qos = QoSProfile(
            reliability=ReliabilityPolicy.BEST_EFFORT,
            history=HistoryPolicy.KEEP_LAST,
            depth=1,
        )
        # Transient local for costmap (Nav2 expects this)
        costmap_qos = QoSProfile(
            reliability=ReliabilityPolicy.RELIABLE,
            history=HistoryPolicy.KEEP_LAST,
            depth=1,
            durability=DurabilityPolicy.TRANSIENT_LOCAL,
        )

        # ── Subscriptions ─────────────────────────────────────────────────────
        self.create_subscription(
            Image, "/camera/color/image_raw", self._image_cb, sensor_qos)

        # ── Publishers ────────────────────────────────────────────────────────
        self._mask_pub  = self.create_publisher(Image, "/segmentation/mask",  10)
        self._cost_pub  = self.create_publisher(
            OccupancyGrid, "/segmentation/costmap", costmap_qos)
        self._debug_pub = self.create_publisher(Image, "/segmentation/debug_image", 1)

        self.get_logger().info(
            f"SidewalkSeg ready  backend={'trt' if isinstance(self._backend, _TRTBackend) else 'onnx'} "
            f"input={self._p['input_width']}x{self._p['input_height']} "
            f"process_every={self._p['process_every_n']} frames"
        )

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

    def _load_params(self) -> dict:
        return {
            "engine_path":        self.get_parameter("engine_path").value,
            "onnx_path":          self.get_parameter("onnx_path").value,
            "input_width":        self.get_parameter("input_width").value,
            "input_height":       self.get_parameter("input_height").value,
            "process_every_n":    self.get_parameter("process_every_n").value,
            "publish_debug":      self.get_parameter("publish_debug_image").value,
            "unknown_as_obstacle":self.get_parameter("unknown_as_obstacle").value,
            "costmap_resolution": self.get_parameter("costmap_resolution").value,
            "costmap_range_m":    self.get_parameter("costmap_range_m").value,
            "camera_height_m":    self.get_parameter("camera_height_m").value,
            "camera_pitch_deg":   self.get_parameter("camera_pitch_deg").value,
        }

    def _init_backend(self):
        p = self._p
        if _TRT_AVAILABLE:
            import os
            if os.path.exists(p["engine_path"]):
                try:
                    self._backend = _TRTBackend(p["engine_path"], self.get_logger())
                    return
                except Exception as e:
                    self.get_logger().warn(f"TRT load failed: {e} — trying ONNX")

        if _ONNX_AVAILABLE:
            import os
            if os.path.exists(p["onnx_path"]):
                try:
                    self._backend = _ONNXBackend(p["onnx_path"], self.get_logger())
                    return
                except Exception as e:
                    self.get_logger().warn(f"ONNX load failed: {e}")

        self.get_logger().warn(
            "No inference backend available — node will publish empty masks. "
            "Run build_engine.py to create the TRT engine."
        )

    # ── Image callback ─────────────────────────────────────────────────────────

    def _image_cb(self, msg: Image):
        self._frame_count += 1
        if self._frame_count % self._p["process_every_n"] != 0:
            return

        try:
            bgr = self._bridge.imgmsg_to_cv2(msg, desired_encoding="bgr8")
        except Exception as e:
            self.get_logger().warn(f"cv_bridge decode error: {e}")
            return

        orig_h, orig_w = bgr.shape[:2]
        iw = self._p["input_width"]
        ih = self._p["input_height"]

        # ── Inference ─────────────────────────────────────────────────────────
        if self._backend is not None:
            try:
                t0 = time.monotonic()
                blob, scale, pad_l, pad_t = preprocess_for_inference(bgr, iw, ih)
                raw = self._backend.infer(blob)
                self._t_infer = time.monotonic() - t0

                # raw shape: (1, num_classes, H, W) or (1, H, W)
                if raw.ndim == 4:
                    class_mask = raw[0].argmax(axis=0).astype(np.uint8)
                else:
                    class_mask = raw[0].astype(np.uint8)

                # Unpad + restore to original resolution
                class_mask_full = unpad_mask(
                    class_mask, orig_h, orig_w, scale, pad_l, pad_t)

                trav_mask = cityscapes_to_traversability(class_mask_full)

            except Exception as e:
                self.get_logger().warn(
                    f"Inference error: {e}", throttle_duration_sec=5.0)
                trav_mask = np.full((orig_h, orig_w), 4, dtype=np.uint8)  # all unknown
        else:
            trav_mask = np.full((orig_h, orig_w), 4, dtype=np.uint8)  # all unknown

        with self._last_mask_lock:
            self._last_mask = trav_mask

        stamp = msg.header.stamp

        # ── Publish segmentation mask ──────────────────────────────────────────
        mask_msg = self._bridge.cv2_to_imgmsg(trav_mask, encoding="mono8")
        mask_msg.header.stamp    = stamp
        mask_msg.header.frame_id = "camera_color_optical_frame"
        self._mask_pub.publish(mask_msg)

        # ── Publish costmap ────────────────────────────────────────────────────
        costmap_msg = self._build_costmap(trav_mask, stamp)
        self._cost_pub.publish(costmap_msg)

        # ── Debug visualisation ────────────────────────────────────────────────
        if self._p["publish_debug"]:
            vis = colorise_traversability(trav_mask)
            debug_msg = self._bridge.cv2_to_imgmsg(vis, encoding="bgr8")
            debug_msg.header.stamp    = stamp
            debug_msg.header.frame_id = "camera_color_optical_frame"
            self._debug_pub.publish(debug_msg)

        if self._frame_count % 30 == 0:
            fps_str = f"{1.0/self._t_infer:.1f} fps" if self._t_infer > 0 else "no inference"
            self.get_logger().debug(
                f"Seg frame {self._frame_count}: {fps_str}",
                throttle_duration_sec=5.0,
            )

    # ── Costmap projection ─────────────────────────────────────────────────────

    def _build_costmap(self, trav_mask: np.ndarray, stamp) -> OccupancyGrid:
        """
        Project the lower portion of the segmentation mask into a flat
        OccupancyGrid in the base_link frame.

        The projection uses a simple inverse-perspective ground model:
          - Only pixels in the lower half of the image are projected
            (upper half is unlikely to be ground plane in front of robot)
          - Each pixel row maps to a forward distance using pinhole geometry
            with the camera mount height and pitch angle.
          - Columns map to lateral (Y) position via horizontal FOV.

        The resulting OccupancyGrid covers [0, costmap_range_m] forward
        and [-costmap_range_m/2, costmap_range_m/2] laterally in base_link.
        """
        p = self._p
        res     = p["costmap_resolution"]    # metres per cell
        rng     = p["costmap_range_m"]       # look-ahead range
        cam_h   = p["camera_height_m"]
        cam_pit = np.deg2rad(p["camera_pitch_deg"])

        # Costmap grid dimensions
        grid_h = int(rng / res)               # forward cells (X direction)
        grid_w = int(rng / res)               # lateral cells (Y direction)
        grid_cx = grid_w // 2                 # lateral centre cell

        cost_map_int8 = traversability_to_costmap(
            trav_mask,
            unknown_as_obstacle=p["unknown_as_obstacle"],
        )

        img_h, img_w = trav_mask.shape
        # Only process lower 60% of image (ground plane region)
        row_start = int(img_h * 0.40)

        # RealSense D435i approximate intrinsics at 640×480
        # (horizontal FOV ~87°, vertical FOV ~58°)
        fx = img_w / (2.0 * np.tan(np.deg2rad(87.0 / 2)))
        fy = img_h / (2.0 * np.tan(np.deg2rad(58.0 / 2)))
        cx = img_w / 2.0
        cy = img_h / 2.0

        # Output occupancy grid (init to -1 = unknown)
        grid = np.full((grid_h, grid_w), -1, dtype=np.int8)

        for row in range(row_start, img_h):
            # Vertical angle from optical axis
            alpha = np.arctan2(-(row - cy), fy) + cam_pit
            if alpha >= 0:
                continue  # ray points up — skip
            # Ground distance from camera base (forward)
            fwd_dist = cam_h / np.tan(-alpha)
            if fwd_dist <= 0 or fwd_dist > rng:
                continue

            grid_row = int(fwd_dist / res)
            if grid_row >= grid_h:
                continue

            for col in range(0, img_w):
                # Lateral angle
                beta = np.arctan2(col - cx, fx)
                lat_dist = fwd_dist * np.tan(beta)

                grid_col = grid_cx + int(lat_dist / res)
                if grid_col < 0 or grid_col >= grid_w:
                    continue

                cell_cost = int(cost_map_int8[row, col])
                existing  = int(grid[grid_row, grid_col])

                # Max-cost merge: most conservative estimate wins
                if existing < 0 or cell_cost > existing:
                    grid[grid_row, grid_col] = np.int8(cell_cost)

        # Build OccupancyGrid message
        msg = OccupancyGrid()
        msg.header.stamp    = stamp
        msg.header.frame_id = "base_link"

        msg.info = MapMetaData()
        msg.info.resolution = res
        msg.info.width      = grid_w
        msg.info.height     = grid_h
        msg.info.map_load_time = stamp

        # Origin: lower-left corner in base_link
        # Grid row 0 = closest (0m forward), row grid_h-1 = rng forward
        # Grid col 0 = leftmost, col grid_cx = straight ahead
        msg.info.origin.position.x = 0.0
        msg.info.origin.position.y = -(rng / 2.0)
        msg.info.origin.position.z = 0.0
        msg.info.origin.orientation.w = 1.0

        msg.data = grid.flatten().tolist()
        return msg


# ── Entry point ───────────────────────────────────────────────────────────────

def main(args=None):
    rclpy.init(args=args)
    node = SidewalkSegNode()
    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    finally:
        node.destroy_node()
        rclpy.try_shutdown()


if __name__ == "__main__":
    main()
