"""uwb_pose_bridge_node.py — UWB→Nav2 localization bridge (Issue #599).

Replaces AMCL by listening to the UWB-IMU EKF fused pose and broadcasting
the map→odom TF that Nav2 requires for global localization.

TF computation (2D, ground robot):
──────────────────────────────────
    map → base_link   := UWB-IMU EKF   (/saltybot/pose/fused_cov)
    odom → base_link  := wheel odometry  (/odom)

    T_map_odom = T_map_base_link  ⊗  T_odom_base_link⁻¹

Published TF:  map → odom  (at UWB rate, ≥ 10 Hz)
Published topic: /initialpose  (once on first valid pose, for Nav2 boot)

Subscribed topics:
    /saltybot/pose/fused_cov  (PoseWithCovarianceStamped) — UWB-IMU EKF
    /odom                     (Odometry)                  — wheel odometry
"""

from __future__ import annotations

import math
from typing import Optional

import rclpy
from rclpy.node import Node
from geometry_msgs.msg import (
    PoseWithCovarianceStamped,
    TransformStamped,
)
from nav_msgs.msg import Odometry
from tf2_ros import TransformBroadcaster


# ── Pure-Python 2-D transform helpers (no numpy / tf_transformations needed) ──

def _yaw_from_q(qx: float, qy: float, qz: float, qw: float) -> float:
    """Extract yaw (rad) from unit quaternion."""
    return math.atan2(2.0 * (qw * qz + qx * qy),
                      1.0 - 2.0 * (qy * qy + qz * qz))


def _q_from_yaw(yaw: float):
    """Return (qx, qy, qz, qw) for a pure-yaw rotation."""
    return 0.0, 0.0, math.sin(yaw * 0.5), math.cos(yaw * 0.5)


def _invert_2d(x: float, y: float, th: float):
    """Invert a 2-D SE(2) transform."""
    c, s = math.cos(th), math.sin(th)
    xi = -(c * x + s * y)
    yi = -(-s * x + c * y)
    return xi, yi, -th


def _compose_2d(x1: float, y1: float, th1: float,
                x2: float, y2: float, th2: float):
    """Compose two 2-D SE(2) transforms:  T1 ⊗ T2."""
    c1, s1 = math.cos(th1), math.sin(th1)
    xr = x1 + c1 * x2 - s1 * y2
    yr = y1 + s1 * x2 + c1 * y2
    thr = th1 + th2
    return xr, yr, thr


def compute_map_odom_tf(
    map_bl_x: float, map_bl_y: float, map_bl_th: float,
    odom_bl_x: float, odom_bl_y: float, odom_bl_th: float,
):
    """Return (x, y, yaw) for the map→odom transform.

    T_map_odom = T_map_base_link ⊗ inv(T_odom_base_link)
    """
    ix, iy, ith = _invert_2d(odom_bl_x, odom_bl_y, odom_bl_th)
    return _compose_2d(map_bl_x, map_bl_y, map_bl_th, ix, iy, ith)


# ── ROS2 node ─────────────────────────────────────────────────────────────────

class UwbPoseBridgeNode(Node):
    """Broadcasts map→odom TF derived from UWB-IMU fused pose + wheel odom.

    Parameters (ROS):
        uwb_pose_topic   (str)   default: /saltybot/pose/fused_cov
        odom_topic       (str)   default: /odom
        tf_publish_hz    (float) default: 20.0   — TF broadcast rate
        tf_timeout_s     (float) default: 0.5    — stop publishing if UWB
                                                   older than this
        map_frame        (str)   default: map
        odom_frame       (str)   default: odom
        base_frame       (str)   default: base_link
    """

    def __init__(self):
        super().__init__("uwb_pose_bridge")

        self.declare_parameter("uwb_pose_topic",  "/saltybot/pose/fused_cov")
        self.declare_parameter("odom_topic",      "/odom")
        self.declare_parameter("tf_publish_hz",   20.0)
        self.declare_parameter("tf_timeout_s",    0.5)
        self.declare_parameter("map_frame",       "map")
        self.declare_parameter("odom_frame",      "odom")
        self.declare_parameter("base_frame",      "base_link")

        self._map_frame  = self.get_parameter("map_frame").value
        self._odom_frame = self.get_parameter("odom_frame").value
        self._base_frame = self.get_parameter("base_frame").value
        self._tf_timeout = self.get_parameter("tf_timeout_s").value

        self._uwb_pose:  Optional[PoseWithCovarianceStamped] = None
        self._odom_pose: Optional[Odometry]                  = None
        self._init_pose_sent: bool = False

        self._tf_broadcaster = TransformBroadcaster(self)

        # /initialpose: Nav2 uses this to seed the robot's initial position
        self._init_pub = self.create_publisher(
            PoseWithCovarianceStamped, "/initialpose", 1)

        self.create_subscription(
            PoseWithCovarianceStamped,
            self.get_parameter("uwb_pose_topic").value,
            self._on_uwb_pose, 10)

        self.create_subscription(
            Odometry,
            self.get_parameter("odom_topic").value,
            self._on_odom, 20)

        hz = self.get_parameter("tf_publish_hz").value
        self.create_timer(1.0 / hz, self._publish_tf)

        self.get_logger().info(
            f"UWB pose bridge ready | "
            f"uwb={self.get_parameter('uwb_pose_topic').value}  "
            f"odom={self.get_parameter('odom_topic').value}  "
            f"tf={self._map_frame}→{self._odom_frame} @ {hz:.0f} Hz"
        )

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

    def _on_uwb_pose(self, msg: PoseWithCovarianceStamped) -> None:
        self._uwb_pose = msg
        if not self._init_pose_sent:
            self._init_pub.publish(msg)
            self._init_pose_sent = True
            self.get_logger().info(
                f"Initial UWB pose published to /initialpose  "
                f"({msg.pose.pose.position.x:.2f}, "
                f"{msg.pose.pose.position.y:.2f})"
            )

    def _on_odom(self, msg: Odometry) -> None:
        self._odom_pose = msg

    # ── TF broadcast ───────────────────────────────────────────────────────

    def _publish_tf(self) -> None:
        if self._uwb_pose is None or self._odom_pose is None:
            return

        # Staleness guard: stop broadcasting if UWB data is too old
        now = self.get_clock().now()
        uwb_stamp = rclpy.time.Time.from_msg(self._uwb_pose.header.stamp)
        age_s = (now - uwb_stamp).nanoseconds * 1e-9
        if age_s > self._tf_timeout:
            self.get_logger().warn(
                f"UWB pose stale ({age_s:.2f}s > {self._tf_timeout}s) — "
                "map→odom TF suppressed",
                throttle_duration_sec=5.0,
            )
            return

        # Extract map→base_link from UWB pose
        p = self._uwb_pose.pose.pose
        map_bl_th = _yaw_from_q(p.orientation.x, p.orientation.y,
                                 p.orientation.z, p.orientation.w)
        map_bl_x, map_bl_y = p.position.x, p.position.y

        # Extract odom→base_link from wheel odometry
        o = self._odom_pose.pose.pose
        odom_bl_th = _yaw_from_q(o.orientation.x, o.orientation.y,
                                  o.orientation.z, o.orientation.w)
        odom_bl_x, odom_bl_y = o.position.x, o.position.y

        # Compute map→odom
        mo_x, mo_y, mo_th = compute_map_odom_tf(
            map_bl_x, map_bl_y, map_bl_th,
            odom_bl_x, odom_bl_y, odom_bl_th,
        )

        # Broadcast
        tf_msg = TransformStamped()
        tf_msg.header.stamp    = now.to_msg()
        tf_msg.header.frame_id = self._map_frame
        tf_msg.child_frame_id  = self._odom_frame
        tf_msg.transform.translation.x = mo_x
        tf_msg.transform.translation.y = mo_y
        tf_msg.transform.translation.z = 0.0
        qx, qy, qz, qw = _q_from_yaw(mo_th)
        tf_msg.transform.rotation.x = qx
        tf_msg.transform.rotation.y = qy
        tf_msg.transform.rotation.z = qz
        tf_msg.transform.rotation.w = qw
        self._tf_broadcaster.sendTransform(tf_msg)


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