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

    def __init__(self, g: float = 9.81):
        self.g = g

    def compute_zmp(self, com_pos, com_acc):
        """计算单点ZMP.

        Args:
            com_pos: 质心位置 [x, y, z]
            com_acc: 质心加速度 [ax, ay, az]

        Returns:
            ZMP位置 [x, y]
        """
        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.size}")
        if com_acc.size < 3:
            raise ValueError(f"com_acc需要3维, 得到{com_acc.size}")
        if com_pos[2] <= 0:
            raise ValueError(f"质心高度必须为正, 得到{com_pos[2]}")

        # 分母保护：避免 (az + g) 接近零导致除零
        denom = com_acc[2] + self.g
        if abs(denom) < 1e-6:
            raise ValueError(f"(az + g) 太小, ZMP无定义")

        # ZMP计算
        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])

def is_in_support_polygon(self, zmp, polygon):
    """判断ZMP是否在支撑多边形内（射线法）."""
    zmp = np.asarray(zmp).flatten()
    polygon = np.asarray(polygon)

    n = len(polygon)
    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) + xi):
            inside = not inside
        j = i
    return inside

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

    def __init__(self, com_height=1.0, g=9.81):
        self.h = com_height
        self.g = g
        self.omega = np.sqrt(g / com_height)  # 自然频率

    def get_natural_freq(self):
        return self.omega

    def get_time_constant(self):
        return 1.0 / self.omega

    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  # 归一化时间 0→1

                # 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))

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

    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):
        self.N = horizon
        self.dt = dt
        self.h = com_height
        self.g = g
        self.mu = friction_coeff
        self.lambda_reg = lambda_reg
        self.kp = kp
        self.kd = kd
        self.omega = np.sqrt(g / com_height)
        self.qp_solver = QPSolver(num_vars=3, horizon=horizon)

    def get_state_matrix(self):
        """获取离散状态空间矩阵 (A, B)."""
        omega = self.omega
        dt = self.dt

        A = np.eye(4)
        A[0, 2] = dt
        A[1, 3] = dt
        A[2, 0] = omega**2 * dt
        A[3, 1] = omega**2 * dt

        B = np.zeros((4, 2))
        B[2, 0] = -omega**2 * dt
        B[3, 1] = -omega**2 * dt

        return A, B

    def robust_estimate(self, target, current, velocity, 
                         disturbance=0.0, eta=0.5):
        """鲁棒控制：PD + 滑模修正项."""
        target = np.asarray(target, dtype=np.float64).flatten()
        current = np.asarray(current, dtype=np.float64).flatten()
        velocity = np.asarray(velocity, dtype=np.float64).flatten()

        # PD基础控制
        error = target - current
        u_pd = self.kp * error - self.kd * velocity

        # 滑模修正项
        lambda_sm = 1.0
        s = error + lambda_sm * velocity
        u_sm = eta * np.sign(s) * abs(disturbance) / 10.0

        return u_pd + u_sm

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

    def run(self, zmp_ref, mpc_controller, disturbance_fn=None):
        """运行闭环仿真."""
        zmp_ref = np.asarray(zmp_ref, dtype=np.float64)
        T = zmp_ref.shape[0]
        omega = self.omega
        dt = self.dt
        h = self.h
        g = self.g
        mass = 70.0

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

            # 当前状态
            if k == 0:
                com_x, com_y = 0.0, 0.0
            else:
                com_x = self.history.com_x[-1]
                com_y = self.history.com_y[-1]

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

            if k >= 2:
                ddcom_x = (dcom_x - prev_dcom_x) / dt
                ddcom_y = (dcom_y - prev_dcom_y) / 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 = c - (h/g) * c̈
            zmp_actual_x = com_x - (h / g) * ddcom_x
            zmp_actual_y = com_y - (h / g) * ddcom_y

            # 摩擦锥检查
            fx, fy = mass * ddcom_x, mass * ddcom_y
            fz = mass * 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

            # 推进到下一时刻
            if k < T - 1:
                new_com_x = com_x + (zmp_ref[k+1, 0] - com_x) * min(0.1, omega * dt)
                new_com_y = com_y + (zmp_ref[k+1, 1] - com_y) * min(0.1, omega * 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 += 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

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

    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)

        # ROS2接口
        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)

        # 100Hz控制循环
        self.timer = self.create_timer(0.01, self._tick_callback)

    def _tick_callback(self):
        """100Hz控制循环."""
        state = self._tick_once()
        # 发布ZMP状态: [t, zmp_x, zmp_y, com_x, com_y, com_z, is_stable]
        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