import rclpy
from rclpy.node import Node
from sensor_msgs.msg import BatteryState, LaserScan
from nav_msgs.msg import Odometry
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
import numpy as np

class RobotMonitor(Node):
    def __init__(self):
        super().__init__('robot_monitor')
        
        # 订阅电池状态
        self.battery_sub = self.create_subscription(
            BatteryState, 'battery_state', self.battery_callback, 10)
        
        # 订阅激光扫描
        self.laser_sub = self.create_subscription(
            LaserScan, 'scan', self.laser_callback, 10)
        
        # 订阅里程计
        self.odom_sub = self.create_subscription(
            Odometry, 'odom', self.odom_callback, 10)
        
        # 初始化数据
        self.battery_voltage = 0.0
        self.laser_ranges = []
        self.robot_x = 0.0
        self.robot_y = 0.0
        
        # 创建图形
        self.fig, (self.ax1, self.ax2, self.ax3) = plt.subplots(3, 1, figsize=(10, 15))
        self.fig.suptitle('Robot Status Monitor')
        
        # 设置电池电量图
        self.battery_bar = self.ax1.bar(['Battery'], [0], color='green')
        self.ax1.set_ylim(0, 100)
        self.ax1.set_ylabel('Battery Level (%)')
        
        # 设置激光扫描图
        self.laser_line, = self.ax2.plot([], [], 'b-')
        self.ax2.set_xlim(-10, 10)
        self.ax2.set_ylim(-10, 10)
        self.ax2.set_aspect('equal')
        self.ax2.set_title('Laser Scan')
        
        # 设置轨迹图
        self.trajectory_line, = self.ax3.plot([], [], 'r-')
        self.robot_marker, = self.ax3.plot([], [], 'bo', markersize=10)
        self.ax3.set_xlim(-5, 5)
        self.ax3.set_ylim(-5, 5)
        self.ax3.set_aspect('equal')
        self.ax3.set_title('Robot Trajectory')
        
        # 轨迹数据
        self.trajectory_x = []
        self.trajectory_y = []
        
        # 创建动画
        self.ani = FuncAnimation(self.fig, self.update_plot, interval=100)
        
        self.get_logger().info('Robot monitor initialized')
    
    def battery_callback(self, msg):
        # 假设电压范围为10.0-12.6V
        voltage = msg.voltage
        self.battery_voltage = max(0.0, min(100.0, (voltage - 10.0) / (12.6 - 10.0) * 100.0))
    
    def laser_callback(self, msg):
        self.laser_ranges = msg.ranges
        self.laser_angles = np.linspace(msg.angle_min, msg.angle_max, len(msg.ranges))
    
    def odom_callback(self, msg):
        self.robot_x = msg.pose.pose.position.x
        self.robot_y = msg.pose.pose.position.y
        self.trajectory_x.append(self.robot_x)
        self.trajectory_y.append(self.robot_y)
        
        # 限制轨迹长度
        if len(self.trajectory_x) > 1000:
            self.trajectory_x = self.trajectory_x[-1000:]
            self.trajectory_y = self.trajectory_y[-1000:]
    
    def update_plot(self, frame):
        # 更新电池图
        self.battery_bar[0].set_height(self.battery_voltage)
        if self.battery_voltage < 20:
            self.battery_bar[0].set_color('red')
        elif self.battery_voltage < 50:
            self.battery_bar[0].set_color('yellow')
        else:
            self.battery_bar[0].set_color('green')
        
        # 更新激光图
        if len(self.laser_ranges) > 0:
            x = np.array(self.laser_ranges) * np.cos(self.laser_angles)
            y = np.array(self.laser_ranges) * np.sin(self.laser_angles)
            self.laser_line.set_data(x, y)
        
        # 更新轨迹图
        self.trajectory_line.set_data(self.trajectory_x, self.trajectory_y)
        self.robot_marker.set_data(self.robot_x, self.robot_y)
        
        # 调整坐标轴范围以适应轨迹
        if len(self.trajectory_x) > 0:
            x_min, x_max = min(self.trajectory_x), max(self.trajectory_x)
            y_min, y_max = min(self.trajectory_y), max(self.trajectory_y)
            margin = 1.0
            self.ax3.set_xlim(min(x_min - margin, -5), max(x_max + margin, 5))
            self.ax3.set_ylim(min(y_min - margin, -5), max(y_max + margin, 5))
        
        return self.battery_bar, self.laser_line, self.trajectory_line, self.robot_marker

def main(args=None):
    rclpy.init(args=args)
    node = RobotMonitor()
    
    # 在单独的线程中运行ROS节点
    import threading
    thread = threading.Thread(target=rclpy.spin, args=(node,), daemon=True)
    thread.start()
    
    # 显示图形
    plt.tight_layout()
    plt.show()
    
    # 关闭节点
    node.destroy_node()
    rclpy.shutdown()
    thread.join()

if __name__ == '__main__':
    main()