#!/usr/bin/env python3
"""
PID 路径跟随 Demo — 最简版本

订阅 /plan (nav_msgs/Path)，收到路径后立即开始 PID 跟随。
不需要 planner_server，不需要状态机，不需要任何额外节点。

用法：
  python3 path_follower_demo.py

然后发一条 /plan 即可：
  ros2 topic pub /plan nav_msgs/msg/Path "{header: {frame_id: 'map'}, poses: [
    {pose: {position: {x: 1.0, y: 0.0}}},
    {pose: {position: {x: 2.0, y: 0.5}}},
    {pose: {position: {x: 3.0, y: 0.0}}}
  ]}"
"""

import rclpy
from rclpy.node import Node
from nav_msgs.msg import Odometry, Path
from geometry_msgs.msg import Twist
import math
import time


class PathFollower(Node):

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

        # ── PID 参数 ──
        self.declare_parameter("kp", 0.8)
        self.declare_parameter("ki", 0.0)
        self.declare_parameter("kd", 0.1)
        self.declare_parameter("linear_speed", 0.3)     # 前进速度 m/s
        self.declare_parameter("waypoint_tolerance", 0.1)  # 到达判定距离 m
        self.declare_parameter("max_angular", 10.0)       # 最大角速度 rad/s

        # ── 订阅 /plan 和 /odom ──
        self.plan_sub = self.create_subscription(
            Path, "/plan", self.plan_cb, 10)
        self.odom_sub = self.create_subscription(
            Odometry, "/odom", self.odom_cb, 10)

        # ── 发布 /cmd_vel ──
        self.cmd_pub = self.create_publisher(Twist, "/cmd_vel", 10)

        # ── 状态 ──
        self.path = []              # [(x, y), ...] 路径点列表(世界坐标)
        self.current_idx = 0        # 当前目标 waypoint 索引
        self.odom = None            # 最新里程计
        self.following = False      # 是否正在跟随

        # ── PID 状态 ──
        self._integral = 0.0
        self._last_err = 0.0
        self._last_time = None

        # ── 主循环 20Hz ──
        self.timer = self.create_timer(0.05, self.control_loop)

        self.get_logger().info("PathFollower ready, waiting for /plan...")

    # ── 回调 ──
    def plan_cb(self, msg: Path):
        """收到新路径，替换当前路径并开始跟随"""
        if len(msg.poses) < 2:
            return

        self.path = [(p.pose.position.x, p.pose.position.y)
                     for p in msg.poses]
        self.current_idx = 0
        self.following = True
        self._reset_pid()
        self.get_logger().info(
            f"Received path with {len(self.path)} waypoints, "
            f"from ({self.path[0][0]:.2f}, {self.path[0][1]:.2f}) "
            f"to ({self.path[-1][0]:.2f}, {self.path[-1][1]:.2f})")

    def odom_cb(self, msg: Odometry):
        self.odom = msg

    # ── PID ──
    def _reset_pid(self):
        self._integral = 0.0
        self._last_err = 0.0
        self._last_time = None

    def _pid(self, error: float, dt: float) -> float:
        kp = self.get_parameter("kp").value
        ki = self.get_parameter("ki").value
        kd = self.get_parameter("kd").value

        self._integral += error * dt
        if dt > 0.001:
            derivative = (error - self._last_err) / dt
        else:
            derivative = 0.0
        self._last_err = error

        return kp * error + ki * self._integral + kd * derivative

    # ── 主控制循环 ──
    def control_loop(self):
        if not self.following or self.odom is None:
            self._publish_cmd(0.0, 0.0)
            return

        # 当前位姿
        x = self.odom.pose.pose.position.x
        y = self.odom.pose.pose.position.y
        yaw = self._quat_to_yaw(self.odom.pose.pose.orientation)

        # 当前目标 waypoint
        if self.current_idx >= len(self.path):
            self._publish_cmd(0.0, 0.0)
            if self.following:
                self.get_logger().info("Path completed!")
                self.following = False
            return

        tx, ty = self.path[self.current_idx]
        dx = tx - x
        dy = ty - y
        dist = math.hypot(dx, dy)

        # 到达当前 waypoint → 进下一个
        tolerance = self.get_parameter("waypoint_tolerance").value
        if dist < tolerance:
            self.current_idx += 1
            if self.current_idx >= len(self.path):
                self._publish_cmd(0.0, 0.0)
                self.get_logger().info("Path completed!")
                self.following = False
                return
            tx, ty = self.path[self.current_idx]
            dx = tx - x
            dy = ty - y

        # 航向误差
        target_heading = math.atan2(dy, dx)
        heading_err = target_heading - yaw
        heading_err = math.atan2(math.sin(heading_err), math.cos(heading_err))

        # PID → 角速度
        now = time.time()
        dt = now - self._last_time if self._last_time else 0.05
        self._last_time = now
        angular = self._pid(heading_err, dt)

        # 限幅
        max_ang = self.get_parameter("max_angular").value
        angular = max(-max_ang, min(max_ang, angular))

        # 速度：误差大时减速
        linear = self.get_parameter("linear_speed").value
        if abs(heading_err) > 0.5:
            linear *= 0.5  # 航向偏差超过 ~30° 时减速
        if abs(heading_err) > 1.0:
            linear *= 0.3  # 偏差超过 ~57° 时更慢

        self._publish_cmd(linear, angular)

        # 每 2 秒打印一次状态
        if int(now * 10) % 20 == 0:
            self.get_logger().info(
                f"WP[{self.current_idx}/{len(self.path)}] "
                f"target=({tx:.2f},{ty:.2f}) dist={dist:.3f} "
                f"err={heading_err:.2f}rad cmd=(v={linear:.2f},ω={angular:.2f})")

    # ── 工具 ──
    def _publish_cmd(self, v, w):
        twist = Twist()
        twist.linear.x = v
        twist.angular.z = w
        self.cmd_pub.publish(twist)

    @staticmethod
    def _quat_to_yaw(q) -> float:
        return math.atan2(2.0 * (q.w * q.z + q.x * q.y),
                          1.0 - 2.0 * (q.y * q.y + q.z * q.z))


def main():
    rclpy.init()
    node = PathFollower()
    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    finally:
        node._publish_cmd(0.0, 0.0)
        node.destroy_node()
        rclpy.shutdown()


if __name__ == "__main__":
    main()