#!/usr/bin/env python

import rospy
import tf2_ros
from geometry_msgs.msg import PoseStamped, PoseWithCovarianceStamped
from nav_msgs.msg import Path, Odometry
from visualization_msgs.msg import Marker
from std_msgs.msg import ColorRGBA
import math

class RobotPoseVisualizer:
    def __init__(self):
        rospy.init_node('robot_pose_visualizer', anonymous=True)
        
        # 发布器
        self.marker_pub = rospy.Publisher('/robot_position_marker', Marker, queue_size=1)
        self.path_pub = rospy.Publisher('/robot_path', Path, queue_size=1)
        self.pose_arrow_pub = rospy.Publisher('/robot_pose_arrow', PoseStamped, queue_size=1)
        self.odom_pub = rospy.Publisher('/robot_odom', Odometry, queue_size=1)
        
        # 订阅器
        self.pose_sub = rospy.Subscriber('/robot_pose', PoseWithCovarianceStamped, 
                                        self.pose_callback, queue_size=10)
        
        # 路径存储
        self.robot_path = Path()
        self.robot_path.header.frame_id = "slamware_map"
        
        # TF broadcaster
        self.tf_broadcaster = tf2_ros.TransformBroadcaster()
        
        # 上一个位置（用于计算距离）
        self.last_position = None
        self.path_distance_threshold = 0.1  # 只有移动超过10cm才添加路径点
        
        rospy.loginfo("🤖 Enhanced Robot Pose Visualizer started")
        rospy.loginfo("📍 Publishing enhanced robot visualization")
        rospy.loginfo("🛤️  Tracking robot trajectory with distance filtering")
        
    def pose_callback(self, msg):
        # 更新时间戳
        current_time = rospy.Time.now()
        
        # 发布机器狗位置标记（圆柱体+箭头）
        self.publish_robot_marker(msg, current_time)
        
        # 更新并发布路径（带距离过滤）
        self.update_robot_path(msg, current_time)
        
        # 发布位姿箭头（大箭头显示当前位置）
        self.publish_pose_arrow(msg, current_time)
        
        # 发布里程计信息（显示轨迹）
        self.publish_odometry(msg, current_time)
        
        # 发布TF变换
        self.publish_tf_transform(msg, current_time)
        
    def publish_robot_marker(self, pose_msg, timestamp):
        """发布机器狗位置标记（红色圆柱体）"""
        marker = Marker()
        marker.header.frame_id = "slamware_map"
        marker.header.stamp = timestamp
        marker.ns = "robot_position"
        marker.id = 0
        marker.type = Marker.CYLINDER
        marker.action = Marker.ADD
        
        # 位置
        marker.pose = pose_msg.pose.pose
        marker.pose.position.z = 0.1  # 稍微抬高一点
        
        # 尺寸（更大更显眼）
        marker.scale.x = 0.4  # 直径
        marker.scale.y = 0.4  # 直径  
        marker.scale.z = 0.3  # 高度
        
        # 颜色（亮红色）
        marker.color.r = 1.0
        marker.color.g = 0.2
        marker.color.b = 0.0
        marker.color.a = 0.9
        
        marker.lifetime = rospy.Duration(2.0)
        
        self.marker_pub.publish(marker)
        
        # 发布机器狗朝向箭头（绿色）
        arrow_marker = Marker()
        arrow_marker.header.frame_id = "slamware_map"
        arrow_marker.header.stamp = timestamp
        arrow_marker.ns = "robot_orientation"
        arrow_marker.id = 1
        arrow_marker.type = Marker.ARROW
        arrow_marker.action = Marker.ADD
        
        # 位置和朝向
        arrow_marker.pose = pose_msg.pose.pose
        arrow_marker.pose.position.z = 0.4  # 箭头高度
        
        # 尺寸（更大的方向箭头）
        arrow_marker.scale.x = 1.0  # 长度
        arrow_marker.scale.y = 0.15  # 宽度
        arrow_marker.scale.z = 0.15  # 高度
        
        # 颜色（亮绿色）
        arrow_marker.color.r = 0.0
        arrow_marker.color.g = 1.0
        arrow_marker.color.b = 0.2
        arrow_marker.color.a = 0.9
        
        arrow_marker.lifetime = rospy.Duration(2.0)
        
        self.marker_pub.publish(arrow_marker)
        
    def update_robot_path(self, pose_msg, timestamp):
        """更新并发布机器狗路径（带距离过滤）"""
        current_pos = pose_msg.pose.pose.position
        
        # 检查是否需要添加新的路径点
        should_add_point = False
        if self.last_position is None:
            should_add_point = True
        else:
            distance = math.sqrt(
                (current_pos.x - self.last_position.x)**2 + 
                (current_pos.y - self.last_position.y)**2
            )
            if distance > self.path_distance_threshold:
                should_add_point = True
        
        if should_add_point:
            # 创建路径点
            pose_stamped = PoseStamped()
            pose_stamped.header.frame_id = "slamware_map"
            pose_stamped.header.stamp = timestamp
            pose_stamped.pose = pose_msg.pose.pose
            
            # 添加到路径（限制路径长度）
            self.robot_path.poses.append(pose_stamped)
            if len(self.robot_path.poses) > 1000:  # 保持最近1000个点
                self.robot_path.poses.pop(0)
                
            self.last_position = current_pos
            
            rospy.loginfo_throttle(5, "🛤️  Path updated: %d points, latest: (%.2f, %.2f)", 
                          len(self.robot_path.poses), current_pos.x, current_pos.y)
            
        # 更新路径时间戳并发布
        self.robot_path.header.stamp = timestamp
        self.path_pub.publish(self.robot_path)
        
    def publish_pose_arrow(self, pose_msg, timestamp):
        """发布大箭头显示当前位置"""
        pose_stamped = PoseStamped()
        pose_stamped.header.frame_id = "slamware_map"
        pose_stamped.header.stamp = timestamp
        pose_stamped.pose = pose_msg.pose.pose
        
        self.pose_arrow_pub.publish(pose_stamped)
        
    def publish_odometry(self, pose_msg, timestamp):
        """发布里程计信息用于轨迹显示"""
        odom = Odometry()
        odom.header.frame_id = "slamware_map"
        odom.header.stamp = timestamp
        odom.child_frame_id = "base_link"
        
        # 位置和朝向
        odom.pose = pose_msg.pose
        
        # 发布里程计
        self.odom_pub.publish(odom)
        
    def publish_tf_transform(self, pose_msg, timestamp):
        """发布TF变换"""
        try:
            import tf2_geometry_msgs
            from geometry_msgs.msg import TransformStamped
            
            transform = TransformStamped()
            transform.header.stamp = timestamp
            transform.header.frame_id = "slamware_map"
            transform.child_frame_id = "base_link"
            
            # 位置
            transform.transform.translation.x = pose_msg.pose.pose.position.x
            transform.transform.translation.y = pose_msg.pose.pose.position.y
            transform.transform.translation.z = pose_msg.pose.pose.position.z
            
            # 朝向
            transform.transform.rotation = pose_msg.pose.pose.orientation
            
            self.tf_broadcaster.sendTransform(transform)
            
        except Exception as e:
            rospy.logwarn_throttle(10, "TF publish failed: %s", str(e))

if __name__ == '__main__':
    try:
        visualizer = RobotPoseVisualizer()
        rospy.spin()
    except rospy.ROSInterruptException:
        rospy.loginfo("Enhanced Robot Pose Visualizer shutting down") 