import numpy as np
from scipy.interpolate import CubicSpline
import matplotlib.pyplot as plt

def generate_spline_points(points, increment):
    """
    使用三次样条曲线拟合并生成根据给定增量生成的拟合点。
    
    :param points: 输入点列表，格式 [(x1, y1), (x2, y2), ...]
    :param increment: x 轴增量
    :return: 拟合点列表 [(x, y), ...]
    """
    # 提取 x 和 y
    x_points, y_points = zip(*points)
    
    # 构建三次样条
    spline = CubicSpline(x_points, y_points)
    
    # 生成根据增量计算的 x 坐标
    x_new = np.arange(min(x_points), max(x_points) + increment, increment)
    
    # 计算 y 坐标
    y_new = spline(x_new)
    
    # 返回拟合点
    return list(zip(x_new, y_new))

if __name__ == "__main__":
    # 输入点
    # 轨迹 1
    # input_points = [(0, 14*3.141592/180), (1, 17*3.141592/180), (2, 2*3.141592/180), (2.5, 0*3.141592/180), (3, 0*3.141592/180)]

    # 轨迹 2 
    input_points = [(0, 16*3.141592/180), (1, 5*3.141592/180), (2, -5*3.141592/180), (3, -16*3.141592/180)]

    # 轨迹 3
    # input_points = [(0, 0), (1, 0), (2, 0), (3, 0)]

    # 轨迹 4 
    # input_points = [(0, -16*3.141592/180), (1, -5*3.141592/180), (2, 5*3.141592/180), (3, 16*3.141592/180)]

    # 轨迹 5
    # input_points = [(0, -14*3.141592/180), (1, -17*3.141592/180), (2, -2*3.141592/180), (2.5, 0*3.141592/180), (3, 0*3.141592/180)]



    increment = 0.1  # 指定 x 轴增量 (将 M 设置为 50，即在 0 到 3 之间，0.06 的增量将产生 51 个点)
    
    # 生成拟合点
    fitted_points = generate_spline_points(input_points, increment)
    
    # 分离 x 和 y 坐标
    x_values = [f"{x:.3f}" for x, y in fitted_points]
    y_values = [f"{y:.3f}" for x, y in fitted_points]
    
    # 打印为 C++ 代码格式
    x_str = "{" + ", ".join(x_values) + "}"
    y_str = "{" + ", ".join(y_values) + "}"
    
    # 打印为 C++ 代码格式
    x_str = "{" + ", ".join(x_values) + "}"
    y_str = "{" + ", ".join(y_values) + "}"
    
    # 将 C++ 代码输出到文件
    with open("src/my_planner/scripts/turning_angle.txt", "w") as f:
        f.write(f"sampling_times.push_back({x_str});\n")
        f.write(f"sampling_turns.push_back({y_str});\n")

    print("拟合点已保存到 src/my_planner/scripts/turning_angle.txt 文件中。")

    # 绘制图像
    plt.figure()  # 创建一个新的图形
    plt.plot([x for x, y in fitted_points], [y for x, y in fitted_points], label='Spline curve', color='b', linewidth=2)
    plt.scatter([x for x, y in input_points], [y for x, y in input_points], color='r', label='Input points', zorder=5)
    
    # 设置标题和标签
    plt.title("Cubic Spline Fitting")
    plt.xlabel("X")
    plt.ylabel("Y")
    
    # 自动调整画布大小
    plt.tight_layout()
    
    # 显示图例
    plt.legend()
    
    # 显示图形
    plt.show()