import numpy as np
from scipy.interpolate import CubicSpline
import math

def generate_spline_trajectory(points, total_time, num_samples=100):
    """
    使用三次样条曲线生成一条平滑轨迹。
    
    :param points: 输入点列表，格式 [(x1, y1), (x2, y2), ...]
    :param total_time: 轨迹总时间
    :param num_samples: 采样点的数量
    :return: 轨迹点列表 [(time, x, y, yaw), ...]
    """
    # 将时间与输入点关联，时间从 0 均匀分布到 total_time
    t_points = np.linspace(0, total_time, len(points))
    x_points, y_points = zip(*points)
    
    # 构建三次样条函数
    cs_x = CubicSpline(t_points, x_points)  # x(t)
    cs_y = CubicSpline(t_points, y_points)  # y(t)
    
    # 采样时间点
    t_samples = np.linspace(0, total_time, num_samples)
    
    # 计算采样点的坐标和角度
    trajectory = []
    for t in t_samples:
        x = cs_x(t)
        y = cs_y(t)
        dx_dt = cs_x.derivative()(t)  # x 的导数
        dy_dt = cs_y.derivative()(t)  # y 的导数
        yaw = math.atan2(dy_dt, dx_dt)  # 计算角度
        
        trajectory.append((t, x, y, yaw))
    
    return trajectory

def generate_xml_trajectory(trajectory):
    """
    将轨迹数据转换为 XML 格式字符串。
    
    :param trajectory: 输入轨迹点列表 [(time, x, y, yaw), ...]
    :return: 格式化的 XML 字符串
    """
    xml_str = ''''''

    # 添加每个 waypoint
    for t, x, y, yaw in trajectory:
        xml_str += f'''               <waypoint>
                  <time>{t:.6f}</time>
                  <pose>{x:.6f} {y:.6f} 0.000000 0.000000 0.000000 {yaw:.6f}</pose>
               </waypoint>\n'''
    return xml_str

if __name__ == "__main__":
    # 输入点：需要连接的 (x, y) 坐标
    input_points = [(-7, -7), (-7, -6), (-7, -5), (-7, -4), (-7, -3), (-7, -2),(-7, -1),(-7, 0),(-7, 1),(-7, 2),(-7, 3),(-7, 4),(-7, 5)]
    total_time = 20.0  # 总时间为 2 秒
    
    # 生成三次样条轨迹
    trajectory = generate_spline_trajectory(input_points, total_time, num_samples=50)
    
    # 生成 XML 格式的轨迹
    xml_output = generate_xml_trajectory(trajectory)
    
    # 保存到文件
    with open("/home/mj/planning_ws/src/turn_on_robot/spline_trajectory.xml", "w") as f:
        f.write(xml_output)
    
    print("轨迹生成完成，结果已保存到 spline_trajectory.xml")