import rclpy
from rclpy.node import Node
from rclpy.qos import QoSProfile, ReliabilityPolicy, HistoryPolicy, DurabilityPolicy
from geometry_msgs.msg import TransformStamped, Point
from visualization_msgs.msg import Marker, MarkerArray
from std_msgs.msg import Float64MultiArray
from builtin_interfaces.msg import Time
import tf2_ros
import numpy as np
from scipy.spatial.transform import Rotation
# 导入第一部分定义的数学类 (假设放在 soft_arm_sim.model 包下)
from soft_arm_sim.model.pcc_kinematics import SoftArmKinematics

class SoftArmSimulator(Node):
    def __init__(self):
        super().__init__('soft_arm_simulator')
        
        # --- 1. 声明并读取 ROS 参数 ---
        # 声明参数名及其默认值。这允许我们在不修改代码的情况下通过 yaml 调整机械臂结构。
        self.declare_parameter('num_sections', 3)          # 分段数
        self.declare_parameter('section_length', 0.240)    # 每段长度
        self.declare_parameter('disks_per_section', 3)     # 每段 Disk 数
        self.declare_parameter('disk_radius', 0.033)       # 物理半径
        self.declare_parameter('visual_disk_radius', 0.04) # 可视化半径 (通常比物理稍大)
        self.declare_parameter('visual_disk_thickness', 0.005) # Disk 厚度
        self.declare_parameter('sim_rate', 30.0)           # 仿真频率 (Hz)

        # 获取参数的实际值
        self.num_sections = self.get_parameter('num_sections').value
        self.section_length = self.get_parameter('section_length').value
        self.disks_per_section = self.get_parameter('disks_per_section').value
        self.disk_radius_kinematics = self.get_parameter('disk_radius').value
        self.visual_disk_r = self.get_parameter('visual_disk_radius').value
        self.visual_disk_h = self.get_parameter('visual_disk_thickness').value
        sim_rate = self.get_parameter('sim_rate').value

        # --- 2. 使用参数初始化数学模型 ---
        # 实例化我们在第一部分编写的运动学类
        self.kinematics = SoftArmKinematics(
            num_sections=self.num_sections,
            section_length=self.section_length,
            disks_per_section=self.disks_per_section,
            disk_radius=self.disk_radius_kinematics
        )
        
        # --- 3. 初始化通信接口 ---
        # TF 广播器：用于发布每个 Disk 的坐标系，让 Rviz 知道它们在哪里
        self.tf_broadcaster = tf2_ros.TransformBroadcaster(self)
        
        # Marker QoS 设置：对于可视化 Marker，使用 RELIABLE 比较稳妥，防止丢包导致模型闪烁
        marker_qos = QoSProfile(
            depth=1,
            reliability=ReliabilityPolicy.RELIABLE,
            history=HistoryPolicy.KEEP_LAST,
            durability=DurabilityPolicy.VOLATILE
        )
        # 发布者：发送 MarkerArray (包含所有 disk 和 骨架线)
        self.marker_pub = self.create_publisher(MarkerArray, 'soft_arm_visual', marker_qos)
        
        # 订阅者：接收控制命令
        # 消息类型 Float64MultiArray，格式预期为 [theta1, phi1, theta2, phi2, ...]
        self.create_subscription(Float64MultiArray, 'soft_arm/command', self.cmd_callback, 10)
        
        # --- 4. 动态初始化 current_config ---
        # current_config 存储当前的机械臂状态。
        # 格式：[(theta, phi, length), (theta, phi, length), ...]
        # 初始状态全为 0 (直立)
        self.current_config = []
        for _ in range(self.num_sections):
            self.current_config.append((0.0, 0.0, self.section_length))
        
        # 创建定时器：以固定频率 (sim_rate) 运行 update_loop
        self.timer = self.create_timer(1.0 / sim_rate, self.update_loop)
        
        self.get_logger().info(f"Soft Arm initialized with {self.num_sections} sections, rate={sim_rate}Hz")

    def cmd_callback(self, msg):
        """
        回调函数：处理收到的控制指令。
        """
        data = msg.data
        # 简单校验：确保数据长度足够覆盖所有段。
        # 因为每段需要 2 个控制量 (弯曲角 theta, 偏转角 phi)
        if len(data) >= 2 * self.num_sections:
            for i in range(self.num_sections):
                idx = i * 2
                # 更新状态配置。注意：此处假设长度 section_length 是不可变的。
                # 如果要做伸缩机器人，可以在这里更新第 3 个参数。
                self.current_config[i] = (data[idx], data[idx+1], self.section_length)

    def update_loop(self):
        """
        主循环：计算运动学 -> 发布 TF -> 发布 Marker
        """
        # 1. 调用数学模型计算所有位置
        # transforms: 所有 Disk 的 4x4 矩阵
        # curve_points: 骨架曲线上的点坐标 (用于画线)
        transforms, curve_points = self.kinematics.forward(self.current_config)
        
        real_time = self.get_clock().now().to_msg()
        
        # 2. 动态生成 disk 名字
        # 这种生成方式 (secX_diskY) 必须与 URDF 或后续的控制逻辑对应。
        disk_names = []
        for s in range(1, self.num_sections + 1):
            for d in range(1, self.disks_per_section + 1):
                disk_names.append(f"sec{s}_disk{d}")

        # 3. 发布 TF 变换
        # 使用 min() 是为了防止数学模型计算出的 transform 数量与预期的 name 数量不一致导致数组越界。
        count = min(len(transforms), len(disk_names))
        
        for i in range(count):
            t = TransformStamped()
            t.header.stamp = real_time
            t.header.frame_id = "base_link"  # 所有变换都是相对于世界/基座的
            t.child_frame_id = disk_names[i] # 子坐标系名字
            
            # 填充位置
            t.transform.translation.x = transforms[i][0, 3]
            t.transform.translation.y = transforms[i][1, 3]
            t.transform.translation.z = transforms[i][2, 3]
            
            # 填充旋转 (将旋转矩阵转换为四元数)
            r = Rotation.from_matrix(transforms[i][:3, :3])
            q = r.as_quat()
            t.transform.rotation.x = q[0]
            t.transform.rotation.y = q[1]
            t.transform.rotation.z = q[2]
            t.transform.rotation.w = q[3]
            
            self.tf_broadcaster.sendTransform(t)

        # 4. 发布可视化 Marker
        self.publish_all_visuals(transforms, curve_points)

    def publish_all_visuals(self, transforms, curve_points):
        """
        辅助函数：构建并发布 MarkerArray 消息，用于在 Rviz 中画出圆盘和线条。
        """
        marker_array = MarkerArray()
        zero_time = Time() # Marker 时间戳通常设为 0，表示“最新”
        
        # --- A. 构建 Disk Markers (圆柱体) ---
        for i, T in enumerate(transforms):
            disk = Marker()
            disk.header.stamp = zero_time
            disk.header.frame_id = "base_link"
            disk.ns = "disks"    # 命名空间
            disk.id = i + 1      # 唯一 ID
            disk.type = Marker.CYLINDER # 形状：圆柱
            disk.action = Marker.ADD    # 动作：添加/修改
            
            # 设置尺寸 (来自参数)
            disk.scale.x = self.visual_disk_r * 2  # 直径 = 半径 * 2
            disk.scale.y = self.visual_disk_r * 2
            disk.scale.z = self.visual_disk_h
            
            disk.color.a = 1.0 # 不透明度
            
            # 颜色逻辑：为了区分不同的 PCC 段，给每段设置不同的颜色 (红/绿/蓝 循环)
            # i // self.disks_per_section 计算当前 disk 属于第几段
            section_idx = i // self.disks_per_section
            
            if section_idx % 3 == 0:
                disk.color.r, disk.color.g, disk.color.b = 1.0, 0.0, 0.0 # 红
            elif section_idx % 3 == 1:
                disk.color.r, disk.color.g, disk.color.b = 0.0, 1.0, 0.0 # 绿
            else:
                disk.color.r, disk.color.g, disk.color.b = 0.0, 0.0, 1.0 # 蓝
            
            # 设置位姿 (直接使用数学模型计算出的矩阵)
            disk.pose.position.x = T[0, 3]
            disk.pose.position.y = T[1, 3]
            disk.pose.position.z = T[2, 3]
            
            r = Rotation.from_matrix(T[:3, :3])
            q = r.as_quat()
            disk.pose.orientation.x = q[0]
            disk.pose.orientation.y = q[1]
            disk.pose.orientation.z = q[2]
            disk.pose.orientation.w = q[3]
            
            # lifetime=0 表示永久显示，直到收到新的更新覆盖它
            disk.lifetime.sec = 0
            disk.lifetime.nanosec = 0
            
            marker_array.markers.append(disk)
            
        # --- B. 构建 Backbone Marker (中心连线) ---
        line = Marker()
        line.header.stamp = zero_time
        line.header.frame_id = "base_link"
        line.ns = "backbone"
        line.id = 0
        line.type = Marker.LINE_STRIP # 形状：线条带
        line.action = Marker.ADD
        line.scale.x = 0.008 # 线条粗细
        line.color.a = 1.0
        line.color.r = 0.9   # 浅灰色
        line.color.g = 0.9
        line.color.b = 0.9
        line.lifetime.sec = 0
        line.lifetime.nanosec = 0
        
        # 填充线条的点
        for p_np in curve_points:
            p = Point()
            p.x, p.y, p.z = float(p_np[0]), float(p_np[1]), float(p_np[2])
            line.points.append(p)
            
        marker_array.markers.append(line)
        
        # 统一发布
        self.marker_pub.publish(marker_array)

def main(args=None):
    rclpy.init(args=args)
    node = SoftArmSimulator()
    rclpy.spin(node) # 保持节点运行，直到按 Ctrl+C
    node.destroy_node()
    rclpy.shutdown()