import numpy as np

def s_curve_jerk_trajectory(
    x0,          # 起始位置
    xg,          # 目标位置
    v_max,       # 最大速度限制
    v0,          # 起始速度
    vg,          # 目标速度
    a_max,       # 最大加速度限制
    j_max,       # 最大加加速度（jerk）限制
    time_step=0.001,  # 时间步长（默认0.001s）
    verbose=False     # 是否打印调试信息
):
    """
    生成带加加速度（jerk）约束的S型轨迹（7段式：加加速→匀加速→减加速→匀速→加减速→匀减速→减减速）
    
    参数:
        x0: 起始位置
        xg: 目标位置
        v_max: 最大允许速度
        v0: 起始速度
        vg: 目标速度
        a_max: 最大允许加速度
        j_max: 最大允许加加速度（jerk）
        time_step: 轨迹采样时间步长（默认0.001s）
        verbose: 是否打印轨迹参数（如各阶段时间、总时间等）
    
    返回:
        t: 时间序列（数组）
        pos: 位置序列（数组）
        vel: 速度序列（数组）
        acc: 加速度序列（数组）
        jerk: 加加速度序列（数组）
    """
    # 初始化变量
    count = 0
    x_g = xg
    x_0 = x0
    v_max = v_max
    v_0 = v0
    v_g = vg
    a_max_input = a_max  # 保存原始a_max用于调整
    
    # --------------------------
    # 1. 计算各阶段时间参数（加速/减速）
    # --------------------------
    # 加速阶段（加加速→匀加速→减加速）
    if (v_max - v_0) * j_max < a_max ** 2:  # 不能达到最大加速度
        if v_0 > v_max:
            Tj1 = 0.0  # 加加速时间
            Ta = 0.0   # 加速总时间
            alima = 0.0  # 实际最大加速度
        else:
            Tj1 = np.sqrt((v_max - v_0) / j_max)  # 加加速时间（对称）
            Ta = 2 * Tj1  # 加速总时间（加加速+减加速，无匀加速）
            alima = Tj1 * j_max  # 实际达到的最大加速度
    else:  # 能达到最大加速度（存在匀加速阶段）
        Tj1 = a_max / j_max  # 加加速时间（达到a_max的时间）
        Ta = Tj1 + (v_max - v_0) / a_max  # 加速总时间（加加速+匀加速+减加速）
        alima = a_max  # 最大加速度为a_max

    # 减速阶段（加减速→匀减速→减减速）
    if (v_max - v_g) * j_max < a_max ** 2:  # 不能达到最大减速度
        Tj2 = np.sqrt((v_max - v_g) / j_max)  # 加减速时间（对称）
        Td = 2 * Tj2  # 减速总时间（加减速+减减速，无匀减速）
        alimd = Tj2 * j_max  # 实际达到的最大减速度
    else:  # 能达到最大减速度（存在匀减速阶段）
        Tj2 = a_max / j_max  # 加减速时间（达到a_max的时间）
        Td = Tj2 + (v_max - v_g) / a_max  # 减速总时间（加减速+匀减速+减减速）
        alimd = a_max  # 最大减速度为a_max

    # --------------------------
    # 2. 计算匀速阶段时间及总时间
    # --------------------------
    # 初始计算匀速时间（可能为负，需调整）
    Tv = (x_g - x_0) / v_max - Ta/2 * (1 + v_0/v_max) - Td/2 * (1 + v_g/v_max)
    T = Tv + Ta + Td  # 初始总时间

    # --------------------------
    # 3. 处理无匀速阶段的情况（Tv ≤ 0）
    # --------------------------
    if Tv > 0:
        # 存在匀速阶段：使用初始计算结果
        vlim = v_max  # 匀速阶段速度
        T = Tv + Ta + Td
    else:
        # 无匀速阶段：重新调整各阶段时间
        Tv = 0.0  # 强制匀速时间为0
        a_max = a_max_input  # 重置a_max
        delta = (a_max ** 4) / (j_max ** 2) + 2 * (v_0 **2 + v_g** 2) + a_max * (4 * (x_g - x_0) - 2 * a_max / j_max * (v_0 + v_g))
        Tj1 = a_max / j_max
        Ta = (a_max **2 / j_max - 2 * v_0 + np.sqrt(delta)) / (2 * a_max)
        Tj2 = a_max / j_max
        Td = (a_max** 2 / j_max - 2 * v_g + np.sqrt(delta)) / (2 * a_max)
        vlim = v_0 + (Ta - Tj1) * alima  # 实际最大速度（未达到v_max）

        # 调整a_max直至Ta、Td合理（避免加速/减速阶段时间过短）
        while Ta < 2 * Tj1 or Td < 2 * Tj2:
            count += 1
            a_max -= a_max_input * 0.1  # 每次减少10%原始a_max
            if a_max <= 0:
                a_max = a_max_input * 0.1  # 避免a_max为负
                break  # 防止无限循环
            
            # 重新计算delta（根据a_max符号调整）
            if a_max > 0:
                delta = (a_max **4) / (j_max** 2) + 2 * (v_0 **2 + v_g** 2) + a_max * (4 * (x_g - x_0) - 2 * a_max / j_max * (v_0 + v_g))
            else:
                delta = (a_max **4) / (j_max** 2) + 2 * (v_0 **2 + v_g** 2) - a_max * (4 * (x_g - x_0) - 2 * a_max / j_max * (v_0 + v_g))
            
            # 重新计算各时间参数
            Tj1 = a_max / j_max if j_max != 0 else 0.0
            Ta = (a_max **2 / j_max - 2 * v_0 + np.sqrt(delta)) / (2 * a_max) if a_max != 0 else 0.0
            Tj2 = a_max / j_max if j_max != 0 else 0.0
            Td = (a_max** 2 / j_max - 2 * v_g + np.sqrt(delta)) / (2 * a_max) if a_max != 0 else 0.0
            vlim = v_0 + (Ta - Tj1) * alima

        # 处理Ta或Td为负的极端情况（速度方向冲突）
        if Ta < 0 or Td < 0:
            if v_0 > v_g:
                # 起始速度 > 目标速度：仅减速
                Ta = 0.0
                Tj1 = 0.0
                alima = 0.0
                Td = 2 * (x_g - x_0) / (v_g + v_0)  # 减速总时间
                # 加减速时间计算
                numerator = j_max * (x_g - x_0) - np.sqrt(j_max * (j_max * (x_g - x_0) **2 + (v_g + v_0)** 2 * (v_g - v_0)))
                denominator = j_max * (v_g + v_0)
                Tj2 = numerator / denominator if denominator != 0 else 0.0
                alimd = -j_max * Tj2  # 实际减速度
                vlim = v_g - (Td - Tj2) * alimd
                alimd = -alimd  # 取绝对值
            else:
                # 起始速度 < 目标速度：仅加速
                Td = 0.0
                Tj2 = 0.0
                alimd = 0.0
                Ta = 2 * (x_g - x_0) / (v_g + v_0)  # 加速总时间
                # 加加速时间计算
                numerator = j_max * (x_g - x_0) - np.sqrt(j_max * (j_max * (x_g - x_0) **2 - (v_g + v_0)** 2 * (v_g - v_0)))
                denominator = j_max * (v_g + v_0)
                Tj1 = numerator / denominator if denominator != 0 else 0.0
                alima = j_max * Tj1  # 实际加速度
                vlim = v_0 + (Ta - Tj1) * alima  # 实际最大速度

        # 更新总时间（无匀速阶段）
        T = Tv + Ta + Td

    # --------------------------
    # 4. 生成轨迹数据（位置、速度、加速度、加加速度）
    # --------------------------
    pos = []  # 位置序列
    vel = []  # 速度序列
    acc = []  # 加速度序列
    jerk = []  # 加加速度序列
    t_total = T  # 总时间
    time_steps = np.arange(0, t_total, time_step)  # 时间序列

    for t in time_steps:
        # 阶段1：加加速（0 ~ Tj1）
        if 0 <= t < Tj1:
            x = x_0 + v_0 * t + j_max * t ** 3 / 6
            v = v_0 + j_max * t ** 2 / 2
            a = j_max * t
            j = j_max
        # 阶段2：匀加速（Tj1 ~ Ta - Tj1）
        elif Tj1 <= t < (Ta - Tj1):
            x = x_0 + v_0 * t + alima / 6 * (3 * t **2 - 3 * Tj1 * t + Tj1** 2)
            v = v_0 + alima * (t - Tj1 / 2)
            a = alima
            j = 0.0
        # 阶段3：减加速（Ta - Tj1 ~ Ta）
        elif (Ta - Tj1) <= t < Ta:
            x = x_0 + (vlim + v_0) * Ta / 2 - vlim * (Ta - t) + j_max * (Ta - t) ** 3 / 6
            v = vlim - j_max * (Ta - t) ** 2 / 2
            a = j_max * (Ta - t)
            j = -j_max
        # 阶段4：匀速（Ta ~ Ta + Tv）
        elif Ta <= t < (Ta + Tv):
            x = x_0 + (vlim + v_0) * Ta / 2 + vlim * (t - Ta)
            v = vlim
            a = 0.0
            j = 0.0
        # 阶段5：加减速（Ta + Tv ~ T - Td + Tj2）即（T - Td ~ T - Td + Tj2）
        elif (T - Td) <= t < (T - Td + Tj2):
            x = x_g - (vlim + v_g) * Td / 2 + vlim * (t - T + Td) - j_max * (t - T + Td) ** 3 / 6
            v = vlim - j_max * (t - T + Td) ** 2 / 2
            a = -j_max * (t - T + Td)
            j = -j_max
        # 阶段6：匀减速（T - Td + Tj2 ~ T - Tj2）
        elif (T - Td + Tj2) <= t < (T - Tj2):
            x = x_g - (vlim + v_g) * Td / 2 + vlim * (t - T + Td) - alimd / 6 * (3 * (t - T + Td) **2 - 3 * Tj2 * (t - T + Td) + Tj2** 2)
            v = vlim - alimd * (t - T + Td - Tj2 / 2)
            a = -alimd
            j = 0.0
        # 阶段7：减减速（T - Tj2 ~ T）
        elif (T - Tj2) <= t < T:
            x = x_g - v_g * (T - t) - j_max * (T - t) ** 3 / 6
            v = v_g + j_max * (T - t) ** 2 / 2
            a = -j_max * (T - t)
            j = j_max
        else:
            # 超出总时间范围（取目标值）
            x = x_g
            v = v_g
            a = 0.0
            j = 0.0

        # 存储数据
        pos.append(x)
        vel.append(v)
        acc.append(a)
        jerk.append(j)

    # --------------------------
    # 5. 打印调试信息（可选）
    # --------------------------
    if verbose:
        print(f"轨迹总时间: {T:.4f} s")
        print(f"加速阶段时间 (Ta): {Ta:.4f} s")
        print(f"减速阶段时间 (Td): {Td:.4f} s")
        print(f"匀速阶段时间 (Tv): {Tv:.4f} s")
        print(f"加加速/加减速时间 (Tj1/Tj2): {Tj1:.4f} s / {Tj2:.4f} s")
        print(f"实际最大速度 (vlim): {vlim:.4f}")
        print(f"实际最大加速度 (alima): {alima:.4f}")
        print(f"实际最大减速度 (alimd): {alimd:.4f}")

    return time_steps, np.array(pos), np.array(vel), np.array(acc), np.array(jerk)


# # --------------------------
# # 示例使用
# # --------------------------
# if __name__ == "__main__":
#     import matplotlib.pyplot as plt

#     # 轨迹参数（可根据需求调整）
#     x0 = 0.0       # 起始位置
#     xg = 60.0      # 目标位置
#     v_max = 20.0   # 最大速度
#     v0 = 10.0      # 起始速度
#     vg = 5.0       # 目标速度
#     a_max = 15.0   # 最大加速度
#     j_max = 20.0   # 最大加加速度
#     time_step = 0.001  # 时间步长

#     # 生成轨迹
#     t, pos, vel, acc, jerk = s_curve_jerk_trajectory(
#         x0=x0, xg=xg,
#         v_max=v_max, v0=v0, vg=vg,
#         a_max=a_max, j_max=j_max,
#         time_step=time_step,
#         verbose=True  # 打印调试信息
#     )

#     # 绘图验证
#     fig, axes = plt.subplots(4, 1, figsize=(10, 16))
#     fig.suptitle("S型加减速轨迹（带加加速度约束）", fontsize=14)

#     # 位置曲线
#     axes[0].plot(t, pos, label="位置")
#     axes[0].axhline(y=xg, color='r', linestyle='--', label="目标位置")
#     axes[0].set_ylabel("位置")
#     axes[0].legend()
#     axes[0].grid(True)

#     # 速度曲线
#     axes[1].plot(t, vel, label="速度")
#     axes[1].axhline(y=v_max, color='r', linestyle='--', label="最大速度")
#     axes[1].set_ylabel("速度")
#     axes[1].legend()
#     axes[1].grid(True)

#     # 加速度曲线
#     axes[2].plot(t, acc, label="加速度")
#     axes[2].axhline(y=a_max, color='r', linestyle='--', label="最大加速度")
#     axes[2].axhline(y=-a_max, color='r', linestyle='--', label="最大减速度")
#     axes[2].set_ylabel("加速度")
#     axes[2].legend()
#     axes[2].grid(True)

#     # 加加速度曲线
#     axes[3].plot(t, jerk, label="加加速度")
#     axes[3].axhline(y=j_max, color='r', linestyle='--', label="最大加加速度")
#     axes[3].axhline(y=-j_max, color='r', linestyle='--', label="最大减加速度")
#     axes[3].set_xlabel("时间 (s)")
#     axes[3].set_ylabel("加加速度")
#     axes[3].legend()
#     axes[3].grid(True)

#     plt.tight_layout()
#     plt.show()
