#!/usr/bin/env python3
"""
birdseye_node — publishes a perspective-corrected top-down view of the floor.

Subscribes to /camera/color/image_raw/compressed (sensor_msgs/CompressedImage).
Publishes  /apollo/birdseye/compressed        (sensor_msgs/CompressedImage).

Parameters (all settable via ROS2 params or env vars at startup):
  camera_height_m        (float, default 1.5)  — camera above floor
  camera_x_offset_m      (float, default 0.0)  — camera ahead of robot center
  camera_y_offset_m      (float, default 0.0)  — camera right of robot center
  camera_pitch_deg       (float, default 0.0)  — physical pitch; 0=level, 90=down

Calibration is loaded from $APOLLO_CAMERA_CALIBRATION_FILE (same env var as
vision_pipeline_node) which must point to a camera_calibration.yaml.

Algorithm
---------
For a camera at height H above a flat floor, pitched downward by a measured
mounting angle, the perspective transform that maps image pixels to floor
coordinates is a planar homography H_img→floor. We compute it analytically
from the camera intrinsics (fx, fy, cx, cy), camera height, and pitch, then
compile it into remap lookup tables.

The output image is centered on the camera ground position. This lets the MOBA
frontend place the sprite at the configured camera offset without requiring a
separate metadata topic for the warped image origin.
"""

import os
import sys
import math
import re
import struct

import rclpy
from rclpy.node import Node
from sensor_msgs.msg import CompressedImage

try:
    import cv2
    import numpy as np
    import yaml
    HAS_DEPS = True
except ImportError as exc:
    HAS_DEPS = False
    _MISSING = str(exc)


# ─── helpers ──────────────────────────────────────────────────────────────────

def _env_float(name: str, default: float) -> float:
    try:
        return float(os.environ.get(name, default))
    except (TypeError, ValueError):
        return default


def _load_calibration(path: str):
    """Return (fx, fy, cx, cy, w, h, dist) from a calibration YAML, or None.

    Handles both the ROS camera_calibration package format (plain YAML) and
    the OpenCV file storage format (%YAML:1.0 header + !!opencv-matrix tags).
    """
    try:
        with open(path, 'r') as f:
            content = f.read()

        # Strip OpenCV-specific YAML constructs that yaml.safe_load can't handle.
        # %YAML:1.0 is not a valid YAML directive (standard is %YAML 1.1).
        content = re.sub(r'^%YAML[^\n]*\n?', '', content, flags=re.MULTILINE)
        # !!opencv-matrix (and any other !! custom tag) → plain mapping
        content = re.sub(r'!![^\s\n]+', '', content)

        data = yaml.safe_load(content)
        if not isinstance(data, dict):
            return None

        # camera_matrix: may be a dict {rows, cols, data} or bare list
        cam = data['camera_matrix']
        m = cam['data'] if isinstance(cam, dict) else cam
        # OpenCV row-major 3×3: [fx, 0, cx, 0, fy, cy, 0, 0, 1]
        fx = float(m[0])
        fy = float(m[4])
        cx = float(m[2])
        cy = float(m[5])

        w = int(data['image_width'])
        h = int(data['image_height'])

        # distortion_coefficients: dict or list
        dist_raw = data.get('distortion_coefficients', {})
        d = dist_raw['data'] if isinstance(dist_raw, dict) else (dist_raw or [])
        dist = np.array(d, dtype=np.float64)

        return fx, fy, cx, cy, w, h, dist
    except Exception:
        return None


def _build_homography(fx: float, fy: float, cx: float, cy: float,
                       img_w: int, img_h: int,
                       camera_height: float,
                       cam_x_offset: float,
                       cam_y_offset: float,
                       camera_pitch_deg: float,
                       output_ppm: float = 200.0):
    """
    Build a homography H that maps image pixels (u, v) to output image pixels.

    Camera pitch is the physical mounting angle measured downward from level:
        0 deg  = optical axis points forward
        90 deg = optical axis points straight down

    The output coordinate system is always top-down. It uses X forward across image width and Y right
    down image height, centered at the camera offset.

    Returns (H_total, out_w, out_h, floor_origin_x, floor_origin_y)
    or None on failure.
    """
    H = camera_height
    if H <= 0 or output_ppm <= 0:
        return None

    pitch = math.radians(max(0.0, min(90.0, camera_pitch_deg)))
    sin_p = math.sin(pitch)
    cos_p = math.cos(pitch)

    def project_to_floor(u: float, v: float):
        x_norm = (u - cx) / fx
        y_norm = (v - cy) / fy

        # OpenCV optical frame ray transformed into robot-local floor coords:
        # X forward, Y right, Z up. Positive image-y points down in the frame.
        ray_x = cos_p - y_norm * sin_p
        ray_y = x_norm
        ray_z = -sin_p - y_norm * cos_p
        if ray_z >= -1e-6:
            return None
        t = -H / ray_z
        return cam_x_offset + t * ray_x, cam_y_offset + t * ray_y

    # Clip rows above horizon. safe_margin bounds the farthest forward floor
    # projection to MAX_OUTPUT_PX pixels at the requested output_ppm.
    #   fwd_distance ≈ H * fy / (safe_margin * cos_p)
    #   → safe_margin = H * fy / (max_floor_m * cos_p)
    MAX_OUTPUT_PX = 1024
    max_floor_m = MAX_OUTPUT_PX / max(output_ppm, 1.0)  # front-half: full budget
    if cos_p > 0.01:
        safe_margin = max(2.0, H * fy / (max_floor_m * cos_p))
    else:
        safe_margin = float(img_h)
    horizon_v = cy - fy * math.tan(pitch)
    top_v = max(0.0, min(float(img_h - 1), horizon_v + safe_margin))

    # Clip rows below v_equator: pixels there project to the floor BEHIND the
    # camera. Only keeping the forward region gives a clean front-half image.
    #   ray_x = 0  when  v = cy + fy * cos_p / sin_p  (= v_equator)
    if sin_p > 0.01:
        v_equator = cy + fy * cos_p / sin_p
    else:
        v_equator = float(img_h)   # level camera: all visible floor is forward
    bot_v = min(float(img_h), v_equator)

    if top_v >= bot_v:
        return None

    src_points = np.array([
        [0.0,          top_v],
        [float(img_w), top_v],
        [float(img_w), bot_v],
        [0.0,          bot_v],
        [cx,           bot_v],
        [cx,           (top_v + bot_v) * 0.5],
    ], dtype=np.float64)

    pairs = []
    for u, v in src_points:
        floor = project_to_floor(float(u), float(v))
        if floor is not None and all(math.isfinite(value) for value in floor):
            pairs.append(((u, v), floor))

    if len(pairs) < 4:
        return None

    floor_points = np.array([floor for _, floor in pairs], dtype=np.float64)
    # Forward-only extent: camera at left edge (x=0), forward goes right.
    fwd_extent = max(0.08, float(np.max(floor_points[:, 0] - cam_x_offset)))
    # Cap lateral extent to MAX_OUTPUT_PX/2 per side so shallow-angle cameras
    # don't produce a portrait image that is 2–3× taller than it is wide.
    # The remap table naturally clips the out-of-bounds lateral region.
    max_lateral_m = (MAX_OUTPUT_PX / 2.0) / max(output_ppm, 1.0)
    max_y_extent = max(0.08, min(max_lateral_m, float(np.max(np.abs(floor_points[:, 1] - cam_y_offset)))))
    out_w = max(16, int(math.ceil(fwd_extent * output_ppm)))
    out_h = max(16, int(math.ceil(max_y_extent * 2.0 * output_ppm)))

    # Camera ground is at left-center of the output image (Pixi anchor x=0, y=0.5).
    dst_points = np.column_stack([
        (floor_points[:, 0] - cam_x_offset) * output_ppm,
        (floor_points[:, 1] - cam_y_offset) * output_ppm + out_h / 2.0,
    ]).astype(np.float32)

    src_points = np.array([src for src, _ in pairs], dtype=np.float32)
    H_mat, _ = cv2.findHomography(src_points, dst_points)
    if H_mat is None:
        return None

    return H_mat, out_w, out_h, cam_x_offset, cam_y_offset - max_y_extent


# ─── ROS node ─────────────────────────────────────────────────────────────────

class BirdseyeNode(Node):

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

        if not HAS_DEPS:
            self.get_logger().error(
                f'birdseye_node: required packages missing — {_MISSING}. Node will not process images.')
            return

        # ── Parameters ──────────────────────────────────────────────────────
        self.camera_height = float(self.declare_parameter(
            'camera_height_m',
            _env_float('APOLLO_BIRDSEYE_HEIGHT_M', 1.5)).value)
        self.cam_x_offset = float(self.declare_parameter(
            'camera_x_offset_m',
            _env_float('APOLLO_BIRDSEYE_X_OFFSET_M',
                       _env_float('APOLLO_TAG_CAMERA_X', 0.0))).value)
        self.cam_y_offset = float(self.declare_parameter(
            'camera_y_offset_m',
            _env_float('APOLLO_BIRDSEYE_Y_OFFSET_M',
                       _env_float('APOLLO_TAG_CAMERA_Y', 0.0))).value)
        self.camera_pitch_deg = float(self.declare_parameter(
            'camera_pitch_deg',
            _env_float('APOLLO_BIRDSEYE_CAMERA_PITCH_DEG',
                       _env_float('APOLLO_BIRDSEYE_ANGLE_DEG', 0.0))).value)
        if self.camera_pitch_deg < 0:
            self.get_logger().warn(
                f'birdseye_node: camera angle {self.camera_pitch_deg:.1f} deg is not usable; '
                'using 0.0 deg.')
            self.camera_pitch_deg = 0.0
        elif self.camera_pitch_deg > 90:
            self.get_logger().warn(
                f'birdseye_node: camera angle {self.camera_pitch_deg:.1f} deg is not usable; '
                'using 90.0 deg.')
            self.camera_pitch_deg = 90.0
        self.output_ppm = float(self.declare_parameter(
            'output_pixels_per_metre',
            _env_float('APOLLO_BIRDSEYE_PPM', 200.0)).value)
        self.max_range_m = float(self.declare_parameter(
            'max_range_m',
            _env_float('APOLLO_BIRDSEYE_MAX_RANGE_M', 0.0)).value)

        # ── Camera intrinsics ────────────────────────────────────────────────
        self._fx = self._fy = self._cx = self._cy = None
        self._img_w = self._img_h = 0
        self._dist_coeffs = None

        calib_path = os.environ.get('APOLLO_CAMERA_CALIBRATION_FILE', '')
        if calib_path:
            result = _load_calibration(calib_path)
            if result:
                self._fx, self._fy, self._cx, self._cy, self._img_w, self._img_h, self._dist_coeffs = result
                self.get_logger().info(
                    f'birdseye_node: loaded calibration from {calib_path} '
                    f'(fx={self._fx:.1f} fy={self._fy:.1f} cx={self._cx:.1f} cy={self._cy:.1f})')
            else:
                self.get_logger().warn(
                    f'birdseye_node: could not parse calibration from {calib_path}; '
                    'using FOV-based defaults.')
        else:
            self.get_logger().warn(
                'birdseye_node: APOLLO_CAMERA_CALIBRATION_FILE not set; '
                'using FOV-based defaults (60° horizontal FOV).')

        self._remap1 = None       # computed on first image if intrinsics known
        self._remap2 = None
        self._warped = None
        self._last_img_shape = None

        # ── Topics ───────────────────────────────────────────────────────────
        self._pub = self.create_publisher(CompressedImage, 'apollo/birdseye/compressed', 5)
        self._sub = self.create_subscription(
            CompressedImage,
            'camera/color/image_raw/compressed',
            self._on_image,
            5,
        )
        self.get_logger().info(
            f'birdseye_node ready — height={self.camera_height:.2f}m '
            f'x_offset={self.cam_x_offset:.3f}m y_offset={self.cam_y_offset:.3f}m '
            f'physical_angle={self.camera_pitch_deg:.1f}deg '
            f'max_range={"%.1fm" % self.max_range_m if self.max_range_m > 0 else "unlimited"} '
            f'ppm={self.output_ppm:.0f}')

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

    def _on_image(self, msg: CompressedImage):
        if not HAS_DEPS:
            return

        # decode
        buf = np.frombuffer(msg.data, dtype=np.uint8)
        frame = cv2.imdecode(buf, cv2.IMREAD_COLOR)
        if frame is None:
            return

        h_img, w_img = frame.shape[:2]

        # Precompute all geometry once per input size; runtime uses only remap.
        if self._remap1 is None or self._last_img_shape != (h_img, w_img):
            self._last_img_shape = (h_img, w_img)
            self._build(w_img, h_img)

        if self._remap1 is None:
            return   # calibration not available

        # Runtime fast path: decode, remap with the precomputed LUT, encode.
        warped = cv2.remap(
            frame,
            self._remap1,
            self._remap2,
            cv2.INTER_LINEAR,
            dst=self._warped,
            borderMode=cv2.BORDER_CONSTANT,
            borderValue=0,
        )

        # encode and publish
        ok, buf_out = cv2.imencode('.jpg', warped, [cv2.IMWRITE_JPEG_QUALITY, 85])
        if not ok:
            return

        out_msg = CompressedImage()
        out_msg.header = msg.header
        out_msg.header.frame_id = 'base_link'
        out_msg.format = 'jpeg'
        out_msg.data = buf_out.tobytes()
        self._pub.publish(out_msg)

    # ── Remap builder ─────────────────────────────────────────────────────────

    def _build(self, w_img: int, h_img: int):
        """Compute the output-to-input remap tables for the current geometry."""
        fx, fy, cx, cy = self._get_intrinsics(w_img, h_img)
        result = _build_homography(
            fx, fy, cx, cy, w_img, h_img,
            self.camera_height, self.cam_x_offset, self.cam_y_offset,
            self.camera_pitch_deg,
            self.output_ppm,
        )
        if result is None:
            self.get_logger().warn('birdseye_node: remap computation failed.')
            self._remap1 = None
            self._remap2 = None
            self._warped = None
            return
        homography, self._out_w, self._out_h, _, _ = result

        # Clip forward extent to the user-configured max range (0 = unlimited).
        if self.max_range_m > 0:
            self._out_w = min(self._out_w, max(16, int(self.max_range_m * self.output_ppm)))

        hfov_half_deg = math.degrees(math.atan2(w_img / 2.0, fx))
        apex = (0, self._out_h // 2)
        fan_r = (int(self.max_range_m * self.output_ppm) if self.max_range_m > 0
                 else int(math.hypot(self._out_w, self._out_h / 2)) + 2)
        n = 64
        pts = [apex]
        for i in range(n + 1):
            a = math.radians(-hfov_half_deg + 2.0 * hfov_half_deg * i / n)
            pts.append((apex[0] + int(fan_r * math.cos(a)),
                        apex[1] + int(fan_r * math.sin(a))))
        valid_mask = np.zeros((self._out_h, self._out_w), dtype=np.uint8)
        cv2.fillPoly(valid_mask, [np.array(pts, dtype=np.int32)], 255)

        inv_h = np.linalg.inv(homography).astype(np.float32)
        x_grid, y_grid = np.meshgrid(
            np.arange(self._out_w, dtype=np.float32),
            np.arange(self._out_h, dtype=np.float32),
        )
        denom = inv_h[2, 0] * x_grid + inv_h[2, 1] * y_grid + inv_h[2, 2]
        safe = np.abs(denom) > 1e-6
        map_x = np.full((self._out_h, self._out_w), -1.0, dtype=np.float32)
        map_y = np.full((self._out_h, self._out_w), -1.0, dtype=np.float32)
        map_x[safe] = (inv_h[0, 0] * x_grid[safe] + inv_h[0, 1] * y_grid[safe] + inv_h[0, 2]) / denom[safe]
        map_y[safe] = (inv_h[1, 0] * x_grid[safe] + inv_h[1, 1] * y_grid[safe] + inv_h[1, 2]) / denom[safe]

        # Fold camera undistortion into the same lookup table when calibration has it.
        if self._dist_coeffs is not None and np.any(self._dist_coeffs != 0):
            k_mat = np.array([
                [fx, 0, cx],
                [0, fy, cy],
                [0, 0, 1],
            ], dtype=np.float64)
            undistort_x, undistort_y = cv2.initUndistortRectifyMap(
                k_mat,
                self._dist_coeffs,
                None,
                k_mat,
                (w_img, h_img),
                cv2.CV_32FC1,
            )
            map_x_distorted = cv2.remap(
                undistort_x,
                map_x,
                map_y,
                cv2.INTER_LINEAR,
                borderMode=cv2.BORDER_CONSTANT,
                borderValue=-1.0,
            )
            map_y_distorted = cv2.remap(
                undistort_y,
                map_x,
                map_y,
                cv2.INTER_LINEAR,
                borderMode=cv2.BORDER_CONSTANT,
                borderValue=-1.0,
            )
            map_x, map_y = map_x_distorted, map_y_distorted

        map_x[valid_mask == 0] = -1.0
        map_y[valid_mask == 0] = -1.0
        self._remap1, self._remap2 = cv2.convertMaps(map_x, map_y, cv2.CV_16SC2)
        self._warped = np.empty((self._out_h, self._out_w, 3), dtype=np.uint8)

    def _get_intrinsics(self, w_img: int, h_img: int):
        if self._fx is not None:
            return self._fx, self._fy, self._cx, self._cy
        # fallback: 84.3° H × 53.6° V FOV (Orbbec Gemini E at 1280×720)
        fx = w_img / (2 * math.tan(math.radians(84.3 / 2)))
        fy = h_img / (2 * math.tan(math.radians(53.6 / 2)))
        cx = w_img / 2.0
        cy = h_img / 2.0
        return fx, fy, cx, cy


# ─── entry point ──────────────────────────────────────────────────────────────

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


if __name__ == '__main__':
    main()
