class ZMPCalculator:
    """ZMP零力矩点计算器."""

    def __init__(self, g: float = 9.81) -> None:
        if g <= 0:
            raise ValueError(f"重力加速度必须为正数, 得到: {g}")
        self.g = g

    def compute_zmp(self, com_pos, com_acc) -> np.ndarray:
        com_pos = np.asarray(com_pos, dtype=np.float64).flatten()
        com_acc = np.asarray(com_acc, dtype=np.float64).flatten()

        if com_pos.size != 3:
            raise ValueError(f"com_pos必须是3维向量, 得到: {com_pos.shape}")
        if com_acc.size != 3:
            raise ValueError(f"com_acc必须是3维向量, 得到: {com_acc.shape}")
        if com_pos[2] < 1e-6:
            raise ValueError(f"质心高度不能为0或负数, 得到: {com_pos[2]}")

        # 分母保护：避免 (az + g) 接近零导致除零
        denom = com_acc[2] + self.g
        if abs(denom) < 1e-9:
            raise ValueError(f"分母异常接近0: com_acc[2] + g = {denom}")

        zmp_x = (com_acc[0] * com_pos[2]) / denom + com_pos[0]
        zmp_y = (com_acc[1] * com_pos[2]) / denom + com_pos[1]
        return np.array([zmp_x, zmp_y], dtype=np.float64)

def is_in_support_polygon(self, zmp, polygon) -> bool:
    zmp = np.asarray(zmp, dtype=np.float64).flatten()
    polygon = np.asarray(polygon, dtype=np.float64)

    n = polygon.shape[0]
    inside = False
    j = n - 1
    for i in range(n):
        xi, yi = polygon[i]
        xj, yj = polygon[j]
        if ((yi > zmp[1]) != (yj > zmp[1])) and \
           (zmp[0] < (xj - xi) * (zmp[1] - yi) / (yj - yi + 1e-12) + xi):
            inside = not inside
        j = i
    return inside

class LIPMTrajectoryGenerator:
    """线性倒立摆轨迹生成器.

    状态空间: x = [c_x, c_y, ċ_x, ċ_y]ᵀ (4维)
    输入: u = [p_x, p_y]ᵀ (ZMP位置)
    动力学: ċ = (g/h)(c - p)
    """

    def __init__(self, com_height: float = 1.0, g: float = 9.81):
        self.h = com_height
        self.g = g
        self.omega = np.sqrt(g / com_height)  # ω = √(g/h)

    def generate_step_trajectory(
        self, step_length, step_period, num_steps, step_width, dt
    ):
        N = int(step_period / dt)
        time_list, com_x_list, com_y_list = [], [], []
        zmp_x_list, zmp_y_list = [], []

        for step in range(num_steps):
            for i in range(N):
                t = i * dt
                tau = t / step_period

                # ZMP分段常数：固定在当前支撑脚中心
                zmp_x = step * step_length + step_length * 0.5
                com_x = step * step_length + step_length * tau

                # Y方向：左右交替
                sign = 1.0 if (step % 4) < 2 else -1.0
                com_y = sign * step_width * abs(tau - 0.5) * 2
                zmp_y = 0.0

                time_list.append(step * step_period + t)
                com_x_list.append(com_x)
                com_y_list.append(com_y)
                zmp_x_list.append(zmp_x)
                zmp_y_list.append(zmp_y)

        return (np.array(time_list), np.array(com_x_list),
                np.array(com_y_list), np.array(zmp_x_list), np.array(zmp_y_list))

def _discretize_lipm(self) -> Tuple[np.ndarray, np.ndarray]:
    """LIPM离散化（前向欧拉）."""
    w2 = self.omega ** 2
    A_c = np.array([
        [0.0, 0.0, 1.0, 0.0],
        [0.0, 0.0, 0.0, 1.0],
        [w2, 0.0, 0.0, 0.0],
        [0.0, w2, 0.0, 0.0],
    ])
    B_c = np.array([
        [0.0, 0.0],
        [0.0, 0.0],
        [-w2, 0.0],
        [0.0, -w2],
    ])
    A_d = np.eye(4) + self.dt * A_c
    B_d = self.dt * B_c
    return A_d, B_d

左脚支撑 ──→ 左脚支撑+右脚摆动 ──→ 右脚支撑 ──→ ...

def bezier_swing_leg(t, p0, p1, p2) -> np.ndarray:
    """三阶贝塞尔曲线.

    B(t) = (1-t)²p₀ + 2(1-t)tp₁ + t²p₂
    t ∈ [0, 1]
    """
    u = 1.0 - t
    return u*u * p0 + 2*u*t * p1 + t*t * p2

def generate_full_gait_cycle(
    step_length=0.3, step_height=0.1, step_period=0.8, dt=0.01
):
    N = int(step_period / dt)
    left_foot = np.zeros((N, 3))
    right_foot = np.zeros((N, 3))

    p_left_ground = np.array([0.0, 0.1, 0.0])
    p_right_ground = np.array([0.0, -0.1, 0.0])

    for i in range(N):
        tau = i / N
        mid_x = p_left_ground[0] + step_length / 2

        if tau < 0.5:  # 前半步: 左脚支撑
            t = tau * 2
            right_foot[i] = bezier_swing_leg(
                t, p_right_ground,
                np.array([mid_x, -0.1, step_height]),
                p_left_ground + np.array([step_length, 0, 0])
            )
            left_foot[i] = p_left_ground
        else:  # 后半步: 右脚支撑
            t = (tau - 0.5) * 2
            left_foot[i] = bezier_swing_leg(
                t, p_left_ground,
                np.array([mid_x, 0.1, step_height]),
                p_right_ground + np.array([step_length, 0, 0])
            )
            right_foot[i] = p_right_ground

    return left_foot, right_foot

def _apply_transition(
    self, com_x_list, zmp_x_list,
    from_gait, to_gait, transition_steps,
    current_x, current_step_length,
):
    if not com_x_list:
        return

    n = min(transition_steps, len(com_x_list[-1]))
    if n <= 1:
        return

    from_step = GAIT_PRESETS[from_gait]["step_length"]
    to_step = GAIT_PRESETS[to_gait]["step_length"]

    last_com_x = com_x_list[-1]
    start_idx = len(last_com_x) - n
    start_val = last_com_x[start_idx]
    end_val = last_com_x[-1]

    for i in range(1, n):
        alpha = i / n
        interp_step = (1 - alpha) * from_step + alpha * to_step
        t = (interp_step * i) / (from_step * n)
        new_val = start_val + (end_val - start_val) * min(t, 1.0)
        # 保证单调递增，防止X回退导致质心反向
        if new_val < last_com_x[start_idx + i - 1]:
            new_val = last_com_x[start_idx + i - 1] + 1e-6
        last_com_x[start_idx + i] = new_val
        last_zmp_x[start_idx + i] = new_val

class MPCController:
    """MPC模型预测控制器."""

    def __init__(
        self, horizon=20, dt=0.01, com_height=1.0,
        g=9.81, friction_coeff=0.5, lambda_reg=0.01,
        kp=10.0, kd=2.0,
    ):
        if horizon < 1:
            raise ValueError(f"预测时域必须为正, 得到: {horizon}")
        # ... 其他参数校验

        self.N = horizon
        self.dt = dt
        self.omega = np.sqrt(g / com_height)
        self.A, self.B = self._discretize_lipm()
        self.qp = QPSolver(num_vars=3, horizon=horizon)

    def predict(self, x0, u_seq):
        """预测状态轨迹."""
        x0 = np.asarray(x0, dtype=np.float64).flatten()
        u_seq = np.asarray(u_seq, dtype=np.float64)
        x_traj = np.zeros((self.N + 1, 4))
        x_traj[0] = x0
        for k in range(self.N):
            x_traj[k + 1] = self.A @ x_traj[k] + self.B @ u_seq[k]
        return x_traj

    def solve(self, zmp_ref):
        """求解MPC优化问题."""
        zmp_ref = np.asarray(zmp_ref, dtype=np.float64)
        return self.qp.solve_zmp_tracking(
            zmp_ref=zmp_ref,
            friction_coeff=self.mu,
            lambda_reg=self.lambda_reg,
        )

def robust_estimate(
    self, target, current, velocity,
    disturbance: float = 0.0, eta: float = 0.5,
) -> np.ndarray:
    """鲁棒控制（PD + 滑模修正项）."""
    target = np.asarray(target).flatten()
    current = np.asarray(current).flatten()
    velocity = np.asarray(velocity).flatten()

    e = target - current
    de = -velocity

    # 滑模面 s = ė + lambda * e
    s = de + 5.0 * e

    u_pd = self.kp * e - self.kd * velocity
    u_sm = eta * np.sign(s) + disturbance

    return u_pd + u_sm

class World:
    """2D步行仿真世界.

    状态: [c_x, c_y, ċ_x, ċ_y]
    动力学: LIPM
    """

    def __init__(self, dt=0.01, com_height=1.0, g=9.81, max_steps=10000):
        self.dt = dt
        self.h = com_height
        self.g = g
        self.omega = np.sqrt(g / com_height)
        self.max_steps = max_steps
        self.state = np.zeros(4)
        self.t = 0.0
        self.history = SimulationResult()
        self.zmp_calc = ZMPCalculator(g=g)

    def run(self, zmp_ref, mpc_controller, disturbance_fn=None):
        zmp_ref = np.asarray(zmp_ref, dtype=np.float64)
        T = zmp_ref.shape[0]

        for k in range(T):
            zmp_ref_x, zmp_ref_y = zmp_ref[k]

            com_x = self.history.com_x[-1] if k > 0 else 0.0
            com_y = self.history.com_y[-1] if k > 0 else 0.0

            # 数值微分求质心加速度
            if k >= 2:
                dcom_x = (com_x - self.history.com_x[-2]) / self.dt
                dcom_y = (com_y - self.history.com_y[-2]) / self.dt
                ddcom_x = (dcom_x - prev_dcom_x) / self.dt
                ddcom_y = (dcom_y - prev_dcom_y) / self.dt
            else:
                ddcom_x = ddcom_y = 0.0

            # 外力扰动
            if disturbance_fn is not None:
                f_dist = np.asarray(disturbance_fn(self.t)).flatten()
                ddcom_x += f_dist[0] / mass
                ddcom_y += f_dist[1] / mass

            # 由LIPM反算ZMP
            zmp_actual_x = com_x - (self.h / self.g) * ddcom_x
            zmp_actual_y = com_y - (self.h / self.g) * ddcom_y

            # 摩擦锥检查
            fx, fy = mass * ddcom_x, mass * ddcom_y
            fz = mass * self.g
            total_fxy = np.sqrt(fx**2 + fy**2)
            friction_ok = total_fxy <= fz * mpc_controller.mu * 1.01

            tracking_err = np.sqrt(
                (zmp_actual_x - zmp_ref_x)**2 +
                (zmp_actual_y - zmp_ref_y)**2
            )
            is_stable = tracking_err < 0.3 and friction_ok

            # 推进：omega*dt系数向ZMP推，避免质心反向
            if k < T - 1:
                new_com_x = com_x + (zmp_ref[k+1, 0] - com_x) * min(0.1, self.omega * self.dt)
                new_com_y = com_y + (zmp_ref[k+1, 1] - com_y) * min(0.1, self.omega * self.dt)

            self.history.com_x.append(new_com_x)
            self.history.com_y.append(new_com_y)
            self.history.zmp_x.append(zmp_ref_x)
            self.history.zmp_y.append(zmp_ref_y)
            self.history.stable.append(is_stable)
            self.t += self.dt

        return self.history

def make_impulse(magnitude=50.0, direction=0.0,
                  start_time=0.5, duration=0.05):
    """脉冲扰动（瞬时推力）."""
    fx_const = magnitude * np.cos(direction)
    fy_const = magnitude * np.sin(direction)

    def disturbance(t):
        if start_time <= t < start_time + duration:
            return np.array([fx_const, fy_const])
        return np.array([0.0, 0.0])
    return disturbance

def make_step(magnitude=30.0, direction=0.0, start_time=0.5):
    """阶跃扰动（从start_time起持续施加）."""
    fx_const = magnitude * np.cos(direction)
    fy_const = magnitude * np.sin(direction)

    def disturbance(t):
        if t >= start_time:
            return np.array([fx_const, fy_const])
        return np.array([0.0, 0.0])
    return disturbance

# 8维状态空间
state = [c_x, c_y, c_z, ċ_x, ċ_y, ċ_z, yaw, ȳaw]
#       ─────────────  ─────────────  ─────────
#       质心位置        质心速度        偏航

def run(self, zmp_ref, mpc_controller, step_period=0.8, yaw_rate_target=0.0):
    c_x, c_y, c_z, dc_x, dc_y, dc_z, yaw, dyaw = self.state

    for k in range(T):
        zmp_ref_x, zmp_ref_y = zmp_ref[k]

        h_target = self._com_height_target(self.t, step_period)

        # 高度PD控制
        kp_z, kd_z = 50.0, 10.0
        ddc_z = kp_z * (h_target - c_z) - kd_z * dc_z

        # 偏航PD控制（加阻尼防飘移）
        kp_yaw, kd_yaw = 5.0, 2.0
        ddyaw = kp_yaw * (yaw_rate_target - dyaw) - kd_yaw * dyaw

        ddc_x = omega**2 * (c_x - zmp_ref_x)
        ddc_y = omega**2 * (c_y - zmp_ref_y)
        zmp_actual_x = c_x - (c_z / g) * ddc_x
        zmp_actual_y = c_y - (c_z / g) * ddc_y

        # omega*dt系数推进，避免质心反向
        step_scale = min(0.1, omega * dt)
        new_c_x = c_x + (zmp_ref[k+1, 0] - c_x) * step_scale
        new_c_y = c_y + (zmp_ref[k+1, 1] - c_y) * step_scale

class ZMPMPCROS2Node(Node):
    """ZMP-MPC ROS2 节点.

    100Hz控制循环，订阅速度指令、发布ZMP状态。
    """

    def __init__(self):
        super().__init__("zmp_mpc_controller")

        self.mpc = MPCController(horizon=20, dt=0.01, com_height=1.0)
        self.lipm = LIPMTrajectoryGenerator(com_height=1.0)
        self.world = World(dt=0.01, com_height=1.0)

        self.zmp_state_pub = self.create_publisher(
            Float32MultiArray, "/zmp_state", 10)
        self.cmd_vel_sub = self.create_subscription(
            Twist, "/cmd_vel", self._cmd_vel_callback, 10)

        self.timer = self.create_timer(0.01, self._tick_callback)

    def _tick_callback(self):
        state = self._tick_once()
        zmp_msg = Float32MultiArray()
        zmp_msg.data = [state.time, state.zmp_x, state.zmp_y,
                        state.com_x, state.com_y, state.com_z,
                        1.0 if state.is_stable else 0.0]
        self.zmp_state_pub.publish(zmp_msg)

# 一键构建
bash scripts/ros2_build.sh

# 启动节点
ros2 run zmp_mpc_controller zmp_mpc_node

# 发送速度指令
ros2 topic pub /cmd_vel geometry_msgs/msg/Twist "{linear: {x: 0.3}}"

# 监听ZMP状态
ros2 topic echo /zmp_state