#!/usr/bin/env python3
"""
initial_pose_to_tf.py — /initialpose → map→odom 的 TF 变换

启动后先发布 identity TF (map 与 odom 重合) 作为回退值。
收到 /initialpose 后查 TF (odom→base_footprint)，反算 map→odom 并更新。

使用: ros2 topic pub /initialpose geometry_msgs/msg/PoseWithCovarianceStamped ...
或在 RViz 中用 2D Pose Estimate 工具设定。
"""

import rclpy
from rclpy.node import Node
from geometry_msgs.msg import PoseWithCovarianceStamped, TransformStamped
from tf2_ros import TransformListener, Buffer, StaticTransformBroadcaster
from tf2_ros import LookupException, ConnectivityException, ExtrapolationException


class InitialPoseToTF(Node):
    def __init__(self):
        super().__init__('initial_pose_to_tf')
        self.declare_parameter('odom_frame', 'odom')
        self.declare_parameter('base_frame', 'base_footprint')
        self.declare_parameter('map_frame', 'map')

        self.odom_frame_ = self.get_parameter('odom_frame').value
        self.base_frame_ = self.get_parameter('base_frame').value
        self.map_frame_ = self.get_parameter('map_frame').value

        self.tf_buffer_ = Buffer()
        self.tf_listener_ = TransformListener(self.tf_buffer_, self)
        self.static_broadcaster_ = StaticTransformBroadcaster(self)

        self.initial_pose_sub_ = self.create_subscription(
            PoseWithCovarianceStamped, '/initialpose',
            self._initial_pose_callback, 10)

        self.get_logger().info(
            f'waiting /initialpose '
            f'(map={self.map_frame_}, odom={self.odom_frame_}, base={self.base_frame_})')

        # fallback identity TF after 1s
        self._identity_timer = self.create_timer(1.0, self._publish_identity)

    def _publish_identity(self):
        self._identity_timer.cancel()
        t = TransformStamped()
        t.header.stamp = self.get_clock().now().to_msg()
        t.header.frame_id = self.map_frame_
        t.child_frame_id = self.odom_frame_
        t.transform.rotation.w = 1.0
        self.static_broadcaster_.sendTransform(t)
        self.get_logger().info('Published identity map→odom TF (fallback)')

    def _initial_pose_callback(self, msg: PoseWithCovarianceStamped):
        try:
            odom_to_base = self.tf_buffer_.lookup_transform(
                self.odom_frame_, self.base_frame_,
                rclpy.time.Time(), rclpy.duration.Duration(seconds=1.0))
        except (LookupException, ConnectivityException, ExtrapolationException) as e:
            self.get_logger().warn(f'TF lookup failed: {e}')
            return

        map_to_base = msg.pose.pose
        ob_t = odom_to_base.transform.translation
        ob_r = odom_to_base.transform.rotation

        # q_base_odom = q_odom_base^(-1)
        q_bo = (-ob_r.x, -ob_r.y, -ob_r.z, ob_r.w)
        # q_map_odom = q_map_base * q_base_odom
        x1, y1, z1, w1 = (
            map_to_base.orientation.x, map_to_base.orientation.y,
            map_to_base.orientation.z, map_to_base.orientation.w)
        x2, y2, z2, w2 = q_bo
        qx = w1 * x2 + x1 * w2 + y1 * z2 - z1 * y2
        qy = w1 * y2 - x1 * z2 + y1 * w2 + z1 * x2
        qz = w1 * z2 + x1 * y2 - y1 * x2 + z1 * w2
        qw = w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2

        # R(q_mo) * p_odom_base
        px, py, pz = ob_t.x, ob_t.y, ob_t.z
        rx = (1 - 2*qy*qy - 2*qz*qz)*px + (2*qx*qy - 2*qz*qw)*py + (2*qx*qz + 2*qy*qw)*pz
        ry = (2*qx*qy + 2*qz*qw)*px + (1 - 2*qx*qx - 2*qz*qz)*py + (2*qy*qz - 2*qx*qw)*pz
        rz = (2*qx*qz - 2*qy*qw)*px + (2*qy*qz + 2*qx*qw)*py + (1 - 2*qx*qx - 2*qy*qy)*pz

        t = TransformStamped()
        t.header.stamp = self.get_clock().now().to_msg()
        t.header.frame_id = self.map_frame_
        t.child_frame_id = self.odom_frame_
        t.transform.translation.x = map_to_base.position.x - rx
        t.transform.translation.y = map_to_base.position.y - ry
        t.transform.translation.z = map_to_base.position.z - rz
        t.transform.rotation.x = qx
        t.transform.rotation.y = qy
        t.transform.rotation.z = qz
        t.transform.rotation.w = qw
        self.static_broadcaster_.sendTransform(t)
        self.get_logger().info(
            f'Updated map→odom: pos=({t.transform.translation.x:.3f},'
            f' {t.transform.translation.y:.3f})')


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


if __name__ == '__main__':
    main()
