"""
visual_odom_node.py — D435i visual odometry at 30 Hz using CUDA optical flow.

Pipeline per frame:
  1. Receive synchronized IR1 (left infrared) + aligned depth images.
  2. Detect/track FAST features with CUDA SparsePyrLKOpticalFlow (GPU path)
     or CPU LK fallback.
  3. Estimate incremental SE(3) via Essential matrix + depth-aided scale.
  4. Publish accumulated pose as nav_msgs/Odometry on /saltybot/visual_odom.

CUDA usage:
  - cv2.cuda.FastFeatureDetector_create   — GPU FAST corner detection
  - cv2.cuda.SparsePyrLKOpticalFlow       — GPU Lucas-Kanade tracking
  Both are ~4× faster than CPU equivalents on the Orin Ampere GPU.
  The node auto-detects CUDA availability and falls back gracefully.

Subscribes:
  /camera/infra1/image_rect_raw     sensor_msgs/Image       30 Hz (grey IR left)
  /camera/depth/image_rect_raw      sensor_msgs/Image       30 Hz float32 depth (m)
  /camera/infra1/camera_info        sensor_msgs/CameraInfo  intrinsics

Publishes:
  /saltybot/visual_odom             nav_msgs/Odometry       30 Hz

Parameters:
  max_features        int    300
  min_features        int    80
  publish_hz          float  30.0   (max rate; actual limited by input)
  frame_id            str    'odom'
  child_frame_id      str    'camera_link'
"""

import time
import math

import rclpy
from rclpy.node import Node
from rclpy.qos import QoSProfile, ReliabilityPolicy, HistoryPolicy, DurabilityPolicy
import message_filters
import numpy as np
import cv2
from cv_bridge import CvBridge

from sensor_msgs.msg import Image, CameraInfo
from nav_msgs.msg import Odometry
from geometry_msgs.msg import (
    Pose, Point, Quaternion,
    Twist, Vector3, TwistWithCovariance, PoseWithCovariance,
)
from std_msgs.msg import Header

from .optical_flow_tracker import OpticalFlowTracker
from .stereo_vo import StereoVO


_SENSOR_QOS = QoSProfile(
    reliability=ReliabilityPolicy.BEST_EFFORT,
    history=HistoryPolicy.KEEP_LAST,
    depth=5,
)
_LATCHED_QOS = QoSProfile(
    reliability=ReliabilityPolicy.RELIABLE,
    history=HistoryPolicy.KEEP_LAST,
    depth=1,
    durability=DurabilityPolicy.TRANSIENT_LOCAL,
)


def _yaw_to_quat(yaw: float) -> Quaternion:
    q = Quaternion()
    q.w = math.cos(yaw * 0.5)
    q.z = math.sin(yaw * 0.5)
    return q


class VisualOdomNode(Node):

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

        self.declare_parameter('max_features',   300)
        self.declare_parameter('min_features',    80)
        self.declare_parameter('frame_id',       'odom')
        self.declare_parameter('child_frame_id', 'camera_link')

        max_feat = self.get_parameter('max_features').value
        min_feat = self.get_parameter('min_features').value
        self._frame_id       = self.get_parameter('frame_id').value
        self._child_frame_id = self.get_parameter('child_frame_id').value

        self._bridge   = CvBridge()
        self._tracker  = OpticalFlowTracker(max_features=max_feat, min_features=min_feat)
        self._vo       = StereoVO()
        self._last_t:  float | None = None

        # Synchronized IR1 + depth
        ir_sub    = message_filters.Subscriber(
            self, Image, '/camera/infra1/image_rect_raw',  qos_profile=_SENSOR_QOS
        )
        depth_sub = message_filters.Subscriber(
            self, Image, '/camera/depth/image_rect_raw',   qos_profile=_SENSOR_QOS
        )
        self._sync = message_filters.ApproximateTimeSynchronizer(
            [ir_sub, depth_sub], queue_size=5, slop=0.033   # 1 frame at 30Hz
        )
        self._sync.registerCallback(self._on_frame)

        self.create_subscription(
            CameraInfo, '/camera/infra1/camera_info', self._on_camera_info, _LATCHED_QOS
        )

        self._odom_pub = self.create_publisher(Odometry, '/saltybot/visual_odom', 10)

        cuda_ok = hasattr(cv2, 'cuda') and cv2.cuda.getCudaEnabledDeviceCount() > 0
        self.get_logger().info(
            f'visual_odom ready — CUDA={"YES" if cuda_ok else "NO (CPU fallback)"}'
        )

    def _on_camera_info(self, msg: CameraInfo) -> None:
        K = np.array(msg.k, dtype=np.float64).reshape(3, 3)
        self._vo.set_K(K)

    def _on_frame(self, ir_msg: Image, depth_msg: Image) -> None:
        now = self.get_clock().now()
        t_sec = now.nanoseconds * 1e-9
        dt    = (t_sec - self._last_t) if self._last_t else (1.0 / 30.0)
        dt    = max(1e-4, min(dt, 0.5))   # clamp to sane range
        self._last_t = t_sec

        # Decode
        try:
            ir_img    = self._bridge.imgmsg_to_cv2(ir_msg,    desired_encoding='mono8')
            depth_img = self._bridge.imgmsg_to_cv2(depth_msg, desired_encoding='32FC1')
        except Exception as exc:
            self.get_logger().warning(f'frame decode error: {exc}')
            return

        # Build grayscale for optical flow (IR is already mono8)
        gray = ir_img

        # Need previous frame — get it from the tracker's internal state
        if self._tracker._prev_gray is None:
            # First frame: just initialise tracker state
            self._tracker._prev_pts  = self._tracker._detect(
                gray, gray.shape[0], gray.shape[1]
            )
            self._tracker._prev_gray = gray
            return

        prev_pts, curr_pts = self._tracker.update(self._tracker._prev_gray, gray)

        est = self._vo.update(prev_pts, curr_pts, depth_img, dt=dt)

        odom = self._build_odom_msg(est, now.to_msg())
        self._odom_pub.publish(odom)

    def _build_odom_msg(self, est, stamp) -> Odometry:
        from .stereo_vo import PoseEstimate
        q = _yaw_to_quat(est.yaw)

        odom = Odometry()
        odom.header.stamp    = stamp
        odom.header.frame_id = self._frame_id
        odom.child_frame_id  = self._child_frame_id

        odom.pose.pose.position    = Point(x=est.x, y=est.y, z=est.z)
        odom.pose.pose.orientation = q

        # Covariance — inflate when VO is invalid
        cov_xy    = 0.05 if est.valid else 0.5
        cov_theta = 0.02 if est.valid else 0.2
        cov_v     = 0.1  if est.valid else 1.0
        p_cov = [0.0] * 36
        p_cov[0]  = cov_xy    # xx
        p_cov[7]  = cov_xy    # yy
        p_cov[35] = cov_theta # yaw-yaw
        odom.pose.covariance = p_cov

        odom.twist.twist.linear.x  = est.vx
        odom.twist.twist.angular.z = est.omega
        t_cov = [0.0] * 36
        t_cov[0]  = cov_v
        t_cov[35] = cov_v * 0.5
        odom.twist.covariance = t_cov

        return odom


def main(args=None):
    rclpy.init(args=args)
    node = VisualOdomNode()
    try:
        rclpy.spin(node)
    finally:
        node.destroy_node()
        rclpy.shutdown()


if __name__ == '__main__':
    main()
