import torch
import torch.nn as nn

class DifferentiablePCC(nn.Module):
    def __init__(self, segment_length=0.24, num_segments=3):
        super().__init__()
        self.register_buffer('L', torch.tensor(segment_length))
        self.n_seg = num_segments

    def forward(self, q):
        """
        q: (Batch, 6)
        Returns:
            tip_pos: (Batch, 3) 末端位置
            backbone_points: (Batch, N_Points, 3) 骨架上所有关键点的坐标
        """
        batch_size = q.shape[0]
        device = q.device
        
        # T 初始化为单位矩阵
        T = torch.eye(4, device=device).unsqueeze(0).repeat(batch_size, 1, 1)
        
        # === 关键修改：我们需要记录原点 (0,0,0) ===
        points = []
        # 先把基座原点放进去 (从 T 中提取位置)
        points.append(T[:, :3, 3]) 

        for i in range(self.n_seg):
            theta = q[:, 2*i]
            phi = q[:, 2*i+1]
            
            T_i = self.pcc_transformation(theta, phi)
            T = torch.bmm(T, T_i)
            
            # 记录每一段的末端位置
            points.append(T[:, :3, 3])

        # 堆叠成 (Batch, N+1, 3)
        backbone_points = torch.stack(points, dim=1)
        
        # 返回末端 和 所有点
        return T[:, :3, 3], backbone_points

    def pcc_transformation(self, theta, phi):
        # ... (保持之前的代码不变，防止 Singular 处理逻辑丢失) ...
        # 如果之前的代码丢了，请告诉我，我再发一遍完整的
        # 这里简略，假设你保留了之前的 pcc_transformation 实现
        mask_straight = torch.abs(theta) < 1e-6
        theta_safe = torch.where(mask_straight, torch.ones_like(theta) * 1e-6, theta)

        c_phi = torch.cos(phi)
        s_phi = torch.sin(phi)
        c_theta = torch.cos(theta_safe)
        s_theta = torch.sin(theta_safe)
        
        dx = self.L * (1 - c_theta) / theta_safe
        dz = self.L * s_theta / theta_safe
        
        dx = torch.where(mask_straight, torch.zeros_like(dx), dx)
        dz = torch.where(mask_straight, self.L, dz)

        batch_size = theta.shape[0]
        T = torch.zeros((batch_size, 4, 4), device=theta.device)
        
        # Rotation
        T[:, 0, 0] = c_phi**2 * (c_theta - 1) + 1
        T[:, 0, 1] = s_phi * c_phi * (c_theta - 1)
        T[:, 0, 2] = c_phi * s_theta
        
        T[:, 1, 0] = s_phi * c_phi * (c_theta - 1)
        T[:, 1, 1] = s_phi**2 * (c_theta - 1) + 1
        T[:, 1, 2] = s_phi * s_theta
        
        T[:, 2, 0] = -c_phi * s_theta
        T[:, 2, 1] = -s_phi * s_theta
        T[:, 2, 2] = c_theta

        # Translation
        T[:, 0, 3] = dx * c_phi
        T[:, 1, 3] = dx * s_phi
        T[:, 2, 3] = dz
        T[:, 3, 3] = 1.0

        return T