"""
surround_vision_node — 360° bird's-eye panoramic view from 4× IMX219 cameras

Each IMX219 image is warped to a top-down (bird's-eye) view using Inverse
Perspective Mapping (IPM). The four patches are composited into a single
overhead image centred on the robot, giving a 360° ground-level view.

Published topics:
  /surround_vision/birdseye  (sensor_msgs/Image, BGR8)
    Square overhead image. Robot is at centre; North (robot +x) is up.
    Scale: configurable, default 0.05 m/pixel → 200×200px = 10 m × 10 m.

Algorithm:
  1. Downsample each image to ipm_width × ipm_height
  2. Compute IPM homography from camera geometry:
       - Source: 4 image points corresponding to known ground distances
       - Destination: corresponding points in bird's-eye canvas
  3. cv2.warpPerspective to produce top-down patch (ipm_patch_size² pixels)
  4. Place each patch in the correct direction on the composite canvas:
       - Front patch → above robot centre
       - Rear patch  → below robot centre
       - Left patch  → left of robot centre
       - Right patch → right of robot centre
  5. Draw robot footprint circle at canvas centre
  6. Publish composite image

Notes:
  - Camera intrinsics are scaled placeholders.  Replace once calibrated.
  - For best results, mount cameras with a slight downward tilt (5-10°).
    Adjust camera_pitch_deg parameter to match hardware.
  - The IPM assumes a flat floor.  Bumps and ramps will cause distortion.
"""

import math

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
from cv_bridge import CvBridge


_CAMERAS = [
    # (name,  yaw_deg,  placement_fn_name)
    ('front', 0.0,    'top'),
    ('left',  90.0,   'left'),
    ('rear',  180.0,  'bottom'),
    ('right', -90.0,  'right'),
]

_CAM_FX_FULL = 229.0
_CAM_FY_FULL = 229.0
_CAM_CX_FULL = 320.0
_CAM_CY_FULL = 240.0


class SurroundVisionNode(Node):
    """Composite 4× IMX219 bird's-eye patches into a 360° overhead image."""

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

        # ── Parameters ──────────────────────────────────────────────────────
        self.declare_parameter('camera_height',   0.30)   # m
        self.declare_parameter('camera_pitch_deg', 0.0)   # tilt down = positive
        self.declare_parameter('ipm_input_width',  160)
        self.declare_parameter('ipm_input_height', 120)
        self.declare_parameter('ipm_patch_px',     80)    # patch size (pixels)
        self.declare_parameter('ipm_range_m',      2.0)   # depth range shown (m)
        self.declare_parameter('canvas_m',         6.0)   # canvas side (m)
        self.declare_parameter('canvas_px',        240)   # canvas pixels per side
        self.declare_parameter('publish_rate',     10.0)  # Hz

        self._bridge = CvBridge()
        self._latest: dict[str, np.ndarray] = {}
        self._H: dict[str, np.ndarray | None] = {n: None for n, *_ in _CAMERAS}

        # ── 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(Image, '/surround_vision/birdseye', 10)

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

        self.get_logger().info(
            'surround_vision: compositing 4× IMX219 → /surround_vision/birdseye'
        )

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

    def _image_cb(self, name: str, msg: Image) -> None:
        try:
            self._latest[name] = self._bridge.imgmsg_to_cv2(msg, 'bgr8')
        except Exception:
            pass

    def _publish(self) -> None:
        canvas_px = self.get_parameter('canvas_px').value
        canvas = np.zeros((canvas_px, canvas_px, 3), dtype=np.uint8)

        cx_px = canvas_px // 2   # robot centre in canvas
        cy_px = canvas_px // 2

        for cam_name, cam_yaw_deg, placement in _CAMERAS:
            if cam_name not in self._latest:
                continue
            patch = self._make_birdseye_patch(cam_name, self._latest[cam_name])
            if patch is None:
                continue
            self._place_patch(canvas, patch, placement, cx_px, cy_px)

        # Draw robot footprint
        canvas_m = self.get_parameter('canvas_m').value
        px_per_m = canvas_px / canvas_m
        robot_radius_px = max(2, int(0.15 * px_per_m))   # robot_radius = 0.15 m
        cv2.circle(canvas, (cx_px, cy_px), robot_radius_px, (0, 255, 0), 1)
        cv2.drawMarker(canvas, (cx_px, cy_px), (0, 255, 0),
                       cv2.MARKER_CROSS, markerSize=6, thickness=1)

        # Draw compass label
        cv2.putText(canvas, 'N', (cx_px - 5, 10),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.4, (200, 200, 200), 1)

        msg = self._bridge.cv2_to_imgmsg(canvas, 'bgr8')
        msg.header.stamp = self.get_clock().now().to_msg()
        msg.header.frame_id = 'base_link'
        self._pub.publish(msg)

    # ── IPM ───────────────────────────────────────────────────────────────────

    def _make_birdseye_patch(self, cam_name: str, img_bgr: np.ndarray) -> np.ndarray | None:
        """Warp one camera image to a top-down bird's-eye patch."""
        iw = self.get_parameter('ipm_input_width').value
        ih = self.get_parameter('ipm_input_height').value
        patch_px = self.get_parameter('ipm_patch_px').value
        h_cam = self.get_parameter('camera_height').value
        pitch_rad = math.radians(self.get_parameter('camera_pitch_deg').value)
        range_m = self.get_parameter('ipm_range_m').value

        small = cv2.resize(img_bgr, (iw, ih), interpolation=cv2.INTER_AREA)

        # Compute / cache homography for this camera at current params
        if self._H[cam_name] is None:
            H = self._compute_ipm_H(iw, ih, h_cam, pitch_rad, range_m, patch_px)
            self._H[cam_name] = H

        if self._H[cam_name] is None:
            return None

        return cv2.warpPerspective(small, self._H[cam_name], (patch_px, patch_px))

    def _compute_ipm_H(
        self, iw: int, ih: int,
        h_cam: float, pitch_rad: float,
        range_m: float, patch_px: int,
    ) -> np.ndarray | None:
        """Compute IPM homography: image plane → bird's-eye top-down patch.

        The IPM maps 4 ground control points, defined by their image coordinates
        (computed from the pinhole model) to their positions in the output patch.

        Output patch layout (for front camera):
          - Top row    = farthest ground (range_m ahead of camera)
          - Bottom row = nearest ground (0.5 m ahead of camera)
          - Left/right = lateral extent ±(range_m/2) m
        """
        fx = _CAM_FX_FULL * iw / 640.0
        fy = _CAM_FY_FULL * ih / 480.0
        cx = _CAM_CX_FULL * iw / 640.0
        cy = _CAM_CY_FULL * ih / 480.0

        d_near = max(0.3, h_cam * 0.5)  # nearest ground distance (avoid singularity)
        d_far = range_m
        lat = range_m / 2.0

        # Image y-coordinate of a ground point at forward distance d:
        # Camera points slightly down by pitch_rad.
        # After pitch rotation, the ground point at distance d is:
        #   y_cam_norm = (h_cam / d - sin(pitch_rad)) / cos(pitch_rad)
        # → v = cy + fy * y_cam_norm
        def _row(d):
            y_n = (h_cam / d - math.sin(pitch_rad)) / math.cos(pitch_rad)
            return cy + fy * y_n

        def _col(d, lat_m):
            return cx + fx * lat_m / d

        v_near = _row(d_near)
        v_far = _row(d_far)
        u_lnear = _col(d_near, -lat)
        u_rnear = _col(d_near,  lat)
        u_lfar  = _col(d_far,  -lat)
        u_rfar  = _col(d_far,   lat)

        # Sanity: far points must be above near points in image
        if v_far >= v_near:
            self.get_logger().warn('IPM: far row >= near row; check camera_height/pitch')
            return None

        src = np.float32([
            [u_rfar,  v_far],   # far right
            [u_lfar,  v_far],   # far left
            [u_rnear, v_near],  # near right
            [u_lnear, v_near],  # near left
        ])
        # In output patch: near = bottom, far = top; right = right, left = left
        p = patch_px - 1
        dst = np.float32([
            [p,   0],   # far right  → top-right
            [0,   0],   # far left   → top-left
            [p,   p],   # near right → bottom-right
            [0,   p],   # near left  → bottom-left
        ])
        return cv2.getPerspectiveTransform(src, dst)

    # ── Canvas placement ──────────────────────────────────────────────────────

    def _place_patch(
        self,
        canvas: np.ndarray,
        patch: np.ndarray,
        placement: str,
        cx_px: int,
        cy_px: int,
    ) -> None:
        """Place a bird's-eye patch on the composite canvas.

        Placement rules (patch near edge = close to robot):
          'top'    → front camera: patch bottom touches robot centre (y grows down)
          'bottom' → rear camera:  patch top touches robot centre
          'left'   → left camera:  patch right edge touches robot centre
          'right'  → right camera: patch left edge touches robot centre
        """
        ph, pw = patch.shape[:2]
        canvas_h, canvas_w = canvas.shape[:2]

        if placement == 'top':       # front: near-ground at bottom of patch → cy_px
            x0, y0 = cx_px - pw // 2, cy_px - ph
        elif placement == 'bottom':  # rear: near-ground at top of patch → cy_px
            # Rotate patch 180° so near-ground faces the robot
            patch = cv2.rotate(patch, cv2.ROTATE_180)
            x0, y0 = cx_px - pw // 2, cy_px
        elif placement == 'left':    # left: near-ground at right edge → cx_px
            # Rotate patch 90° CW so near-ground faces right (toward robot)
            patch = cv2.rotate(patch, cv2.ROTATE_90_CLOCKWISE)
            ph, pw = patch.shape[:2]
            x0, y0 = cx_px - pw, cy_px - ph // 2
        elif placement == 'right':   # right: near-ground at left edge → cx_px
            # Rotate patch 90° CCW so near-ground faces left (toward robot)
            patch = cv2.rotate(patch, cv2.ROTATE_90_COUNTERCLOCKWISE)
            ph, pw = patch.shape[:2]
            x0, y0 = cx_px, cy_px - ph // 2
        else:
            return

        # Clip to canvas bounds
        x1, y1 = x0 + pw, y0 + ph
        px0 = max(0, -x0);     py0 = max(0, -y0)
        cx0 = max(0,  x0);     cy0 = max(0,  y0)
        cx1 = min(canvas_w, x1); cy1 = min(canvas_h, y1)
        if cx1 <= cx0 or cy1 <= cy0:
            return

        canvas[cy0:cy1, cx0:cx1] = patch[py0:py0 + (cy1 - cy0), px0:px0 + (cx1 - cx0)]


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


if __name__ == '__main__':
    main()
