import numpy as np

class PCCSection:
    """
    PCCSection 类：表示柔性臂的单个独立分段 (Segment)。
    每个分段由若干个 Disk (圆盘) 组成，且假设该段内部曲率恒定。
    """
    def __init__(self, length, disk_num, disk_radius):
        self.L0 = length           # 该段的总长度 (弧长)
        self.n_disks = disk_num    # 该段包含的 Disk 数量
        self.r = disk_radius       # Disk 半径 (用于物理计算，此处暂未深度使用)
        self.d_per_segment = length / disk_num  # 两个 Disk 之间的沿弧长距离

    def _get_transform_at_s(self, s, q):
        """
        核心函数：计算沿 PCC 曲线 s 处的齐次变换矩阵 T (4x4)。
        
        参数:
            s: 沿曲线的弧长位置 (0 <= s <= L0)
            q: 构型空间坐标 [theta, phi, s_total]
               theta: 弯曲角度 (决定曲率 k = theta / s_total)
               phi:   弯曲方向 (绕 Z 轴的偏转角)
               s_total: 当前段的实际弧长 (通常等于 L0，除非考虑伸缩)
        """
        theta, phi, s_total = q
        T = np.eye(4)  # 初始化单位矩阵
        
        # --- 奇异点处理 ---
        # 当 theta 接近 0 时，机械臂处于直立状态。
        # 此时曲率 k -> 0，半径 R -> 无穷大，直接套用 PCC 公式会导致除以零错误。
        # 因此，直立状态下直接简化为沿 Z 轴的平移。
        if abs(theta) < 1e-6:
            T[0, 3] = 0
            T[1, 3] = 0
            T[2, 3] = s  # 纯粹沿 Z 轴向上延伸 s 长度
        else:
            # --- PCC 变换公式 ---
            # 1. 曲率 k = theta / s_total
            # 2. 变换逻辑：先旋转 phi 到弯曲平面 -> 在平面内弯曲 theta -> 旋转 -phi 回去 (等效推导结果如下)
            
            k = theta / s_total
            
            # 预计算三角函数，减少重复计算
            ct = np.cos(theta * (s / s_total))  # 当前位置的切线角度余弦
            st = np.sin(theta * (s / s_total))  # 当前位置的切线角度正弦
            cp = np.cos(phi)                    # 弯曲平面的方位角余弦
            sp = np.sin(phi)                    # 弯曲平面的方位角正弦
            
            # 旋转矩阵 R (3x3)
            # 描述了当前截面相对于底部的旋转姿态
            R = np.array([
                [cp*cp*(ct-1)+1, cp*sp*(ct-1), cp*st],
                [cp*sp*(ct-1),   sp*sp*(ct-1)+1, sp*st],
                [-cp*st,         -sp*st,         ct]
            ])
            
            # 位置向量 p (3x1)
            # 描述了当前截面中心在空间中的坐标 (x, y, z)
            p = np.array([
                (cp * (1 - ct)) / k,  # x 坐标：在弯曲方向上的投影
                (sp * (1 - ct)) / k,  # y 坐标
                st / k                # z 坐标：高度
            ])
            
            T[:3, :3] = R
            T[:3, 3] = p
            
        return T

    def forward_kinematics(self, q):
        """
        计算该段内所有 Disk 相对于该段底部的变换矩阵。
        用于放置可视化模型 (Cylinder)。
        """
        theta, phi, s_total = q
        transforms = []
        # 遍历该段内的每一个 Disk (索引从1开始，不包括底座)
        for i in range(1, self.n_disks + 1):
            # 计算第 i 个 disk 在曲线上的弧长位置
            s = (s_total / self.n_disks) * i
            # 调用核心公式计算变换
            transforms.append(self._get_transform_at_s(s, q))
        return transforms

    def get_curve_points(self, q, num_points=10):
        """
        生成用于可视化的骨架曲线点集 (Marker LineStrip)。
        比 forward_kinematics 采样更密集，使线条看起来更平滑。
        """
        theta, phi, s_total = q
        points = []
        # 在 0 到 s_total 之间均匀生成 num_points 个点
        s_values = np.linspace(0, s_total, num_points)
        for s in s_values:
            T = self._get_transform_at_s(s, q)
            points.append(T[:3, 3]) # 只提取位置 (x, y, z)
        return points

class SoftArmKinematics:
    """
    SoftArmKinematics 类：多段 PCC 机械臂的高层管理器。
    负责将多个 PCCSection 串联起来，计算全局坐标。
    """
    def __init__(self, num_sections=3, section_length=0.24, disks_per_section=3, disk_radius=0.033):
        """
        初始化时接收动态参数，支持任意段数和长度的配置。
        """
        self.num_sections = num_sections
        self.sections = []
        
        # 根据参数动态生成 Section 对象列表
        for _ in range(num_sections):
            self.sections.append(PCCSection(
                length=section_length,
                disk_num=disks_per_section,
                disk_radius=disk_radius
            ))

    def forward(self, joint_configs):
        """
        计算整个机械臂的全局运动学。
        
        参数:
            joint_configs: 列表，包含每一段的配置 [(theta1, phi1, L1), (theta2, phi2, L2), ...]
            
        返回:
            all_transforms: 所有 Disk 相对于世界坐标系 (Base) 的变换矩阵列表。
            all_path_points: 整个机械臂骨架曲线的全局坐标点列表。
        """
        all_transforms = []
        all_path_points = []
        
        # T_current_base: 当前段的基坐标系。
        # 初始为单位矩阵 (机械臂根部在世界原点)
        T_current_base = np.eye(4)
        all_path_points.append(T_current_base[:3, 3]) # 添加原点
        
        # 安全检查：防止输入的配置数量少于定义的段数 (例如只给了1组数据但有3段)
        limit = min(len(joint_configs), len(self.sections))
        
        for i in range(limit):
            config = joint_configs[i]
            
            # 1. 计算当前段的局部变换 (相对于当前段的底部)
            section_transforms = self.sections[i].forward_kinematics(config)
            local_points = self.sections[i].get_curve_points(config, num_points=15)
            
            # 2. 坐标系转换：局部 -> 全局
            # 公式：P_global = T_current_base * P_local
            for p_local in local_points:
                p_global = T_current_base[:3, :3] @ p_local + T_current_base[:3, 3]
                all_path_points.append(p_global)
            
            for T_local in section_transforms:
                # 矩阵乘法级联变换
                T_global = T_current_base @ T_local
                all_transforms.append(T_global)
            
            # 3. 更新基坐标系
            # 下一段的底部 = 当前段的末端 (即当前段 transform 列表中的最后一个)
            if section_transforms:
                T_current_base = all_transforms[-1]
                
        return all_transforms, all_path_points