"""
surround_costmap_node — Camera-based obstacle detection for Nav2 costmap

Subscribes to 4× IMX219 CSI camera feeds, projects potential obstacles onto
the ground plane using a pinhole model + known camera geometry, and publishes
them as a PointCloud2 on /surround_vision/obstacles.

Nav2's obstacle_layer subscribes to this PointCloud2 as an observation source,
supplementing RPLIDAR with obstacles that LIDAR misses:
  - Transparent glass walls / doors
  - Low furniture (chair legs, coffee tables)
  - People detected by camera but below LIDAR scan plane

Algorithm (V1):
  1. Downsample each image to processing resolution (default 160×120)
  2. Undistort with placeholder IMX219 fisheye calibration
     (replace K, D with real calibration data once measured)
  3. Detect edges (Canny) — edges indicate object boundaries
  4. For each edge pixel below the horizon (v > cy), back-project to ground
     using pinhole model + camera height:
       distance = h_cam * fy / (row - cy)
       lateral  = (col - cx) * distance / fx
  5. Filter by range [min_obstacle_range, max_obstacle_range]
  6. Rotate from camera frame to base_link using known camera yaw
  7. Publish merged point cloud at obstacle height (z ≈ 0.5 m above ground)

Camera yaws in base_link frame:
  front=0°, left=90°, rear=180°, right=-90°

Extension path (Phase 2c+):
  Replace Canny with TensorRT person/obstacle detector for higher precision.
  Use optical flow to distinguish moving obstacles from static background.
"""

import math
import struct

import cv2
import numpy as np

import rclpy
from rclpy.node import Node
from rclpy.qos import QoSProfile, ReliabilityPolicy, HistoryPolicy

from sensor_msgs.msg import Image, PointCloud2, PointField
from cv_bridge import CvBridge


# ── Camera geometry ──────────────────────────────────────────────────────────
# Yaw angle (degrees) of each camera's optical axis in base_link frame.
# 0° = facing +x (front), 90° = facing +y (left).
_CAMERAS = [
    ('front', 0.0),
    ('left',  90.0),
    ('rear',  180.0),
    ('right', -90.0),
]

# IMX219 160° fisheye placeholder calibration (640×480 resolution).
# fx/fy computed for equidistant fisheye: f = r_max / (FOV/2 in radians)
# FOV=160°, r_max=320px → f = 320 / 1.396 ≈ 229 px
# After undistortion to rectilinear, effective FOV ≈ 100°.
# For the simplified pinhole model post-undistortion at 160×120:
# fx = fy ≈ (80 / tan(50°)) ≈ 67 px
_CAM_FX_FULL = 229.0
_CAM_FY_FULL = 229.0
_CAM_CX_FULL = 320.0
_CAM_CY_FULL = 240.0
# Fisheye distortion coefficients k1..k4 (cv2.fisheye model).
# Zero means no correction beyond the equidistant model assumption.
# Replace with real values from `ros2 run camera_calibration cameracalibrator`.
_CAM_D = np.array([0.0, 0.0, 0.0, 0.0], dtype=np.float64)


class SurroundCostmapNode(Node):
    """Publish PointCloud2 of camera-detected obstacles for Nav2 costmap."""

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

        # ── Parameters ──────────────────────────────────────────────────────
        self.declare_parameter('camera_height',       0.30)   # m above ground
        self.declare_parameter('camera_pitch_deg',    0.0)    # tilt down = positive
        self.declare_parameter('proc_width',          160)    # processing resolution
        self.declare_parameter('proc_height',         120)
        self.declare_parameter('min_obstacle_range',  0.25)   # m from camera
        self.declare_parameter('max_obstacle_range',  3.00)   # m from camera
        self.declare_parameter('obstacle_pub_height', 0.50)   # z in PointCloud2 (m)
        self.declare_parameter('canny_low',           40)
        self.declare_parameter('canny_high',          120)
        self.declare_parameter('publish_rate',        5.0)    # Hz

        self._bridge = CvBridge()
        self._latest: dict[str, Image] = {}

        # ── Subscriptions ────────────────────────────────────────────────────
        best_effort = QoSProfile(
            reliability=ReliabilityPolicy.BEST_EFFORT,
            history=HistoryPolicy.KEEP_LAST,
            depth=1,
        )
        for name, _ in _CAMERAS:
            self.create_subscription(
                Image,
                f'/camera/{name}/image_raw',
                lambda msg, n=name: self._image_cb(n, msg),
                best_effort,
            )

        # ── Publisher ────────────────────────────────────────────────────────
        self._pub = self.create_publisher(PointCloud2, '/surround_vision/obstacles', 10)

        # ── Timer ────────────────────────────────────────────────────────────
        rate = self.get_parameter('publish_rate').value
        self.create_timer(1.0 / rate, self._publish)

        self.get_logger().info(
            'surround_costmap: listening on /camera/{front,left,rear,right}/image_raw'
        )

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

    def _image_cb(self, name: str, msg: Image) -> None:
        self._latest[name] = msg

    def _publish(self) -> None:
        all_pts: list[tuple[float, float, float]] = []
        stamp = self.get_clock().now().to_msg()

        for cam_name, cam_yaw_deg in _CAMERAS:
            if cam_name not in self._latest:
                continue
            try:
                img_bgr = self._bridge.imgmsg_to_cv2(self._latest[cam_name], 'bgr8')
            except Exception as exc:
                self.get_logger().warn(f'{cam_name}: cv_bridge error: {exc}', throttle_duration_sec=5.0)
                continue

            pts = self._process_camera(img_bgr, cam_yaw_deg)
            all_pts.extend(pts)

        if not all_pts:
            return

        self._pub.publish(self._make_pc2(all_pts, 'base_link', stamp))

    # ── Core processing ───────────────────────────────────────────────────────

    def _process_camera(
        self, img_bgr: np.ndarray, cam_yaw_deg: float
    ) -> list[tuple[float, float, float]]:
        """Detect edge-based obstacles and project to base_link ground coords."""

        pw = self.get_parameter('proc_width').value
        ph = self.get_parameter('proc_height').value
        h_cam = self.get_parameter('camera_height').value
        pitch_rad = math.radians(self.get_parameter('camera_pitch_deg').value)
        d_min = self.get_parameter('min_obstacle_range').value
        d_max = self.get_parameter('max_obstacle_range').value
        obs_z = self.get_parameter('obstacle_pub_height').value
        canny_lo = self.get_parameter('canny_low').value
        canny_hi = self.get_parameter('canny_high').value

        # ── Downsample & fisheye undistort ───────────────────────────────────
        small = cv2.resize(img_bgr, (pw, ph), interpolation=cv2.INTER_AREA)
        small = self._undistort(small)

        # ── Effective intrinsics at processing resolution ────────────────────
        scale_x = pw / 640.0
        scale_y = ph / 480.0
        fx = _CAM_FX_FULL * scale_x
        fy = _CAM_FY_FULL * scale_y
        cx = _CAM_CX_FULL * scale_x
        cy = _CAM_CY_FULL * scale_y

        # ── Obstacle detection via Canny edge detection ──────────────────────
        gray = cv2.cvtColor(small, cv2.COLOR_BGR2GRAY)
        # Equalise histogram for consistent thresholding across lighting conditions
        gray = cv2.equalizeHist(gray)
        edges = cv2.Canny(gray, canny_lo, canny_hi)

        # Morphological close: connect nearby edge fragments
        kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
        edges = cv2.morphologyEx(edges, cv2.MORPH_CLOSE, kernel)

        # ── Ground projection ─────────────────────────────────────────────────
        # Only pixels below the horizon (v > cy adjusted for pitch) can project
        # onto the ground plane.
        # Effective horizon row accounting for camera pitch:
        #   pitch > 0 → camera tilts down → horizon shifts up (smaller v_horizon)
        v_horizon = cy - fy * math.tan(pitch_rad)

        rows, cols = np.where(edges > 0)
        valid_mask = rows > v_horizon
        rows = rows[valid_mask].astype(np.float64)
        cols = cols[valid_mask].astype(np.float64)

        if len(rows) == 0:
            return []

        # Back-project to camera frame (z forward, x right, y down)
        # Camera at height h_cam, ground plane at y_cam = h_cam (y is down)
        # For pitch angle φ (camera tilted down), adjust the horizon:
        # Normalized ray direction: [x_n, y_n, 1]
        x_n = (cols - cx) / fx
        y_n = (rows - cy) / fy

        # Rotate ray by camera pitch (about x-axis, positive = tilting down)
        cos_p = math.cos(pitch_rad)
        sin_p = math.sin(pitch_rad)
        y_cam = y_n * cos_p + sin_p       # y component in camera frame
        z_cam = -y_n * sin_p + cos_p      # z component in camera frame (forward)

        # Ground intersection: camera y = h_cam → ray param t = h_cam / y_cam
        valid_ground = y_cam > 0.01       # must be pointing toward floor
        y_cam = y_cam[valid_ground]
        z_cam = z_cam[valid_ground]
        x_n_vg = x_n[valid_ground]

        t = h_cam / y_cam
        x_ground_cam = x_n_vg * t        # lateral in camera frame
        z_ground_cam = z_cam * t          # forward in camera frame

        # ── Range filter ─────────────────────────────────────────────────────
        dist = np.hypot(x_ground_cam, z_ground_cam)
        in_range = (dist >= d_min) & (dist <= d_max)
        x_g = x_ground_cam[in_range]
        z_g = z_ground_cam[in_range]

        if len(x_g) == 0:
            return []

        # ── Rotate from camera frame to base_link ────────────────────────────
        # Camera coordinate convention: z forward, x right.
        # Robot base_link:              x forward, y left.
        # Camera yaw θ = angle of camera +z axis from base_link +x axis.
        #
        # base_x =  cos(θ) * z_g  -  sin(θ) * x_g
        # base_y =  sin(θ) * z_g  +  cos(θ) * x_g
        # base_z =  obs_z  (obstacle height in base_link)
        yaw = math.radians(cam_yaw_deg)
        cos_y = math.cos(yaw)
        sin_y = math.sin(yaw)
        base_x = cos_y * z_g - sin_y * x_g
        base_y = sin_y * z_g + cos_y * x_g

        # Subsample to limit point cloud density (every 4th point)
        step = max(1, len(base_x) // 200)
        pts = [
            (float(bx), float(by), obs_z)
            for bx, by in zip(base_x[::step], base_y[::step])
        ]
        return pts

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

    def _undistort(self, img: np.ndarray) -> np.ndarray:
        """Approximate fisheye undistortion with placeholder IMX219 calibration.

        The fisheye calibration matrices are scaled to the processing resolution.
        Replace _CAM_D with values from camera_calibration once measured.
        """
        h, w = img.shape[:2]
        K = np.array([
            [_CAM_FX_FULL * w / 640.0, 0.0, _CAM_CX_FULL * w / 640.0],
            [0.0, _CAM_FY_FULL * h / 480.0, _CAM_CY_FULL * h / 480.0],
            [0.0, 0.0, 1.0],
        ], dtype=np.float64)
        # cv2.fisheye.undistortImage handles equidistant fisheye model.
        # With D=[0,0,0,0] this is a no-op (use as hook for real calibration).
        if np.any(_CAM_D != 0.0):
            return cv2.fisheye.undistortImage(img, K, _CAM_D, Knew=K)
        return img

    @staticmethod
    def _make_pc2(
        pts: list[tuple[float, float, float]],
        frame_id: str,
        stamp,
    ) -> PointCloud2:
        """Pack a list of (x, y, z) tuples into a PointCloud2 message."""
        fields = [
            PointField(name='x', offset=0,  datatype=PointField.FLOAT32, count=1),
            PointField(name='y', offset=4,  datatype=PointField.FLOAT32, count=1),
            PointField(name='z', offset=8,  datatype=PointField.FLOAT32, count=1),
        ]
        data = struct.pack(f'{3 * len(pts)}f', *[v for pt in pts for v in pt])
        msg = PointCloud2()
        msg.header.stamp = stamp
        msg.header.frame_id = frame_id
        msg.height = 1
        msg.width = len(pts)
        msg.fields = fields
        msg.is_bigendian = False
        msg.point_step = 12
        msg.row_step = 12 * len(pts)
        msg.is_dense = True
        msg.data = data
        return msg


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


if __name__ == '__main__':
    main()
