"""
midas_depth_node.py -- MiDaS monocular depth estimation for saltybot.

Uses MiDaS_small via ONNX Runtime or TensorRT FP16.
Provides relative depth estimates for cameras without active depth (CSI cameras).

Publishes /social/depth/cam{i}/image  (sensor_msgs/Image, float32, relative depth)
"""

import os
import numpy as np
import rclpy
from rclpy.node import Node
from rclpy.qos import QoSProfile, ReliabilityPolicy, HistoryPolicy
from sensor_msgs.msg import Image
from cv_bridge import CvBridge

# MiDaS_small input size
_MIDAS_H = 256
_MIDAS_W = 256
# ImageNet normalization
_MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32)
_STD = np.array([0.229, 0.224, 0.225], dtype=np.float32)


class _TRTBackend:
    """TensorRT inference backend for MiDaS."""

    def __init__(self, engine_path: str, logger):
        self._logger = logger
        try:
            import tensorrt as trt
            import pycuda.driver as cuda
            import pycuda.autoinit  # noqa: F401

            self._cuda = cuda
            rt_logger = trt.Logger(trt.Logger.WARNING)
            with open(engine_path, 'rb') as f:
                engine = trt.Runtime(rt_logger).deserialize_cuda_engine(f.read())
            self._context = engine.create_execution_context()

            # Allocate buffers
            self._d_input = cuda.mem_alloc(1 * 3 * _MIDAS_H * _MIDAS_W * 4)
            self._d_output = cuda.mem_alloc(1 * _MIDAS_H * _MIDAS_W * 4)
            self._h_output = np.empty((_MIDAS_H, _MIDAS_W), dtype=np.float32)
            self._stream = cuda.Stream()
            self._logger.info(f'TRT engine loaded: {engine_path}')
        except Exception as e:
            raise RuntimeError(f'TRT init failed: {e}')

    def infer(self, input_tensor: np.ndarray) -> np.ndarray:
        self._cuda.memcpy_htod_async(
            self._d_input, input_tensor.ravel(), self._stream)
        self._context.execute_async_v2(
            bindings=[int(self._d_input), int(self._d_output)],
            stream_handle=self._stream.handle)
        self._cuda.memcpy_dtoh_async(
            self._h_output, self._d_output, self._stream)
        self._stream.synchronize()
        return self._h_output.copy()


class _ONNXBackend:
    """ONNX Runtime inference backend for MiDaS."""

    def __init__(self, onnx_path: str, logger):
        self._logger = logger
        try:
            import onnxruntime as ort
            providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
            self._session = ort.InferenceSession(onnx_path, providers=providers)
            self._input_name = self._session.get_inputs()[0].name
            self._logger.info(f'ONNX model loaded: {onnx_path}')
        except Exception as e:
            raise RuntimeError(f'ONNX init failed: {e}')

    def infer(self, input_tensor: np.ndarray) -> np.ndarray:
        result = self._session.run(None, {self._input_name: input_tensor})
        return result[0].squeeze()


class MiDaSDepthNode(Node):
    """MiDaS monocular depth estimation node."""

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

        # Parameters
        self.declare_parameter('onnx_path',
                               '/mnt/nvme/saltybot/models/midas_small.onnx')
        self.declare_parameter('engine_path',
                               '/mnt/nvme/saltybot/models/midas_small.engine')
        self.declare_parameter('camera_topics', [
            '/surround/cam0/image_raw',
            '/surround/cam1/image_raw',
            '/surround/cam2/image_raw',
            '/surround/cam3/image_raw',
        ])
        self.declare_parameter('output_scale', 1.0)
        self.declare_parameter('process_rate', 5.0)

        onnx_path = self.get_parameter('onnx_path').value
        engine_path = self.get_parameter('engine_path').value
        self._camera_topics = self.get_parameter('camera_topics').value
        self._output_scale = self.get_parameter('output_scale').value
        process_rate = self.get_parameter('process_rate').value

        # Initialize inference backend (TRT preferred, ONNX fallback)
        self._backend = None
        if os.path.exists(engine_path):
            try:
                self._backend = _TRTBackend(engine_path, self.get_logger())
            except RuntimeError:
                self.get_logger().warn('TRT failed, trying ONNX fallback')
        if self._backend is None and os.path.exists(onnx_path):
            try:
                self._backend = _ONNXBackend(onnx_path, self.get_logger())
            except RuntimeError:
                self.get_logger().error('Both TRT and ONNX backends failed')
        if self._backend is None:
            self.get_logger().error(
                'No MiDaS model found. Depth estimation disabled.')

        self._bridge = CvBridge()

        # Latest frames per camera (round-robin processing)
        self._latest_frames = [None] * len(self._camera_topics)
        self._current_cam_idx = 0

        # QoS for camera subscriptions
        cam_qos = QoSProfile(
            reliability=ReliabilityPolicy.BEST_EFFORT,
            history=HistoryPolicy.KEEP_LAST,
            depth=1,
        )

        # Subscribe to each camera topic
        self._cam_subs = []
        for i, topic in enumerate(self._camera_topics):
            sub = self.create_subscription(
                Image, topic,
                lambda msg, idx=i: self._on_image(msg, idx),
                cam_qos)
            self._cam_subs.append(sub)

        # Publishers: one per camera
        self._depth_pubs = []
        for i in range(len(self._camera_topics)):
            pub = self.create_publisher(
                Image, f'/social/depth/cam{i}/image', 10)
            self._depth_pubs.append(pub)

        # Timer: round-robin across cameras
        timer_period = 1.0 / process_rate
        self._timer = self.create_timer(timer_period, self._timer_callback)

        self.get_logger().info(
            f'MiDaS depth node started: {len(self._camera_topics)} cameras '
            f'@ {process_rate} Hz')

    def _on_image(self, msg: Image, cam_idx: int):
        """Cache latest frame for each camera."""
        self._latest_frames[cam_idx] = msg

    def _preprocess(self, bgr: np.ndarray) -> np.ndarray:
        """Preprocess BGR image to MiDaS input tensor [1,3,256,256]."""
        import cv2
        rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB)
        resized = cv2.resize(rgb, (_MIDAS_W, _MIDAS_H),
                             interpolation=cv2.INTER_LINEAR)
        normalized = (resized.astype(np.float32) / 255.0 - _MEAN) / _STD
        # HWC -> CHW -> NCHW
        tensor = normalized.transpose(2, 0, 1)[np.newaxis, ...]
        return tensor.astype(np.float32)

    def _infer(self, tensor: np.ndarray) -> np.ndarray:
        """Run inference, returns [256,256] float32 relative inverse depth."""
        if self._backend is None:
            return np.zeros((_MIDAS_H, _MIDAS_W), dtype=np.float32)
        return self._backend.infer(tensor)

    def _postprocess(self, raw: np.ndarray, orig_shape: tuple) -> np.ndarray:
        """Resize depth back to original image shape, apply output_scale."""
        import cv2
        h, w = orig_shape[:2]
        depth = cv2.resize(raw, (w, h), interpolation=cv2.INTER_LINEAR)
        depth = depth * self._output_scale
        return depth

    def _timer_callback(self):
        """Process one camera per tick (round-robin)."""
        if not self._camera_topics:
            return

        idx = self._current_cam_idx
        self._current_cam_idx = (idx + 1) % len(self._camera_topics)

        msg = self._latest_frames[idx]
        if msg is None:
            return

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

        tensor = self._preprocess(bgr)
        raw_depth = self._infer(tensor)
        depth_map = self._postprocess(raw_depth, bgr.shape)

        # Publish as float32 Image
        depth_msg = self._bridge.cv2_to_imgmsg(depth_map, encoding='32FC1')
        depth_msg.header = msg.header
        self._depth_pubs[idx].publish(depth_msg)


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


if __name__ == '__main__':
    main()
