#!/usr/bin/env python3
"""Pure pursuit path follower for Nav2 autonomous navigation.

Implements the pure pursuit algorithm for following a path or trajectory.
The algorithm computes steering commands to make the robot follow a path
by targeting a lookahead point on the path.

The pure pursuit algorithm:
1. Finds the closest point on the path to the robot
2. Looks ahead a specified distance along the path
3. Computes a circular arc passing through robot position to lookahead point
4. Publishes velocity commands to follow this arc
"""

import math
import rclpy
from rclpy.node import Node
from nav_msgs.msg import Path, Odometry
from geometry_msgs.msg import Twist, PoseStamped
from std_msgs.msg import Float32


class PurePursuitNode(Node):
    """ROS2 node implementing pure pursuit path follower."""

    def __init__(self):
        super().__init__('pure_pursuit')

        # Parameters
        self.declare_parameter('lookahead_distance', 0.5)
        self.declare_parameter('goal_tolerance', 0.1)
        self.declare_parameter('heading_tolerance', 0.1)
        self.declare_parameter('max_linear_velocity', 1.0)
        self.declare_parameter('max_angular_velocity', 1.57)
        self.declare_parameter('linear_velocity_scale', 1.0)
        self.declare_parameter('angular_velocity_scale', 1.0)
        self.declare_parameter('use_heading_correction', True)
        self.declare_parameter('publish_frequency', 10)
        self.declare_parameter('enable_path_following', True)

        # Read parameters
        self.lookahead_distance = self.get_parameter('lookahead_distance').value
        self.goal_tolerance = self.get_parameter('goal_tolerance').value
        self.heading_tolerance = self.get_parameter('heading_tolerance').value
        self.max_linear_velocity = self.get_parameter('max_linear_velocity').value
        self.max_angular_velocity = self.get_parameter('max_angular_velocity').value
        self.linear_velocity_scale = self.get_parameter('linear_velocity_scale').value
        self.angular_velocity_scale = self.get_parameter('angular_velocity_scale').value
        self.use_heading_correction = self.get_parameter('use_heading_correction').value
        publish_frequency = self.get_parameter('publish_frequency').value
        self.enable_path_following = self.get_parameter('enable_path_following').value

        # Current state
        self.current_pose = None
        self.current_path = None
        self.goal_reached = False

        # Subscriptions
        self.sub_odom = self.create_subscription(
            Odometry, '/odom', self._on_odometry, 10
        )
        self.sub_path = self.create_subscription(
            Path, '/path', self._on_path, 10
        )

        # Publishers
        self.pub_cmd_vel = self.create_publisher(Twist, '/cmd_vel_tracked', 10)
        self.pub_tracking_error = self.create_publisher(Float32, '/tracking_error', 10)

        # Timer for control loop
        period = 1.0 / publish_frequency
        self.timer = self.create_timer(period, self._control_loop)

        self.get_logger().info(
            f"Pure pursuit initialized. "
            f"Lookahead: {self.lookahead_distance}m, "
            f"Goal tolerance: {self.goal_tolerance}m"
        )

    def _on_odometry(self, msg: Odometry) -> None:
        """Callback for odometry messages."""
        self.current_pose = msg.pose.pose

    def _on_path(self, msg: Path) -> None:
        """Callback for path messages."""
        self.current_path = msg

    def _quaternion_to_yaw(self, quat):
        """Convert quaternion to yaw angle."""
        siny_cosp = 2 * (quat.w * quat.z + quat.x * quat.y)
        cosy_cosp = 1 - 2 * (quat.y * quat.y + quat.z * quat.z)
        yaw = math.atan2(siny_cosp, cosy_cosp)
        return yaw

    def _yaw_to_quaternion(self, yaw):
        """Convert yaw angle to quaternion."""
        from geometry_msgs.msg import Quaternion
        cy = math.cos(yaw * 0.5)
        sy = math.sin(yaw * 0.5)
        return Quaternion(x=0.0, y=0.0, z=sy, w=cy)

    def _distance(self, p1, p2):
        """Compute Euclidean distance between two points."""
        dx = p1.x - p2.x
        dy = p1.y - p2.y
        return math.sqrt(dx * dx + dy * dy)

    def _find_closest_point_on_path(self):
        """Find closest point on path to current robot position."""
        if self.current_path is None or len(self.current_path.poses) < 2:
            return None, 0, float('inf')

        closest_idx = 0
        closest_dist = float('inf')

        # Find closest waypoint
        for i, pose in enumerate(self.current_path.poses):
            dist = self._distance(self.current_pose.position, pose.pose.position)
            if dist < closest_dist:
                closest_dist = dist
                closest_idx = i

        return closest_idx, closest_dist, closest_idx

    def _find_lookahead_point(self):
        """Find lookahead point on path ahead of closest point."""
        if self.current_path is None or len(self.current_path.poses) < 2:
            return None

        closest_idx, _, _ = self._find_closest_point_on_path()
        if closest_idx is None:
            return None

        # Find point at lookahead distance along path
        current_dist = 0.0
        for i in range(closest_idx, len(self.current_path.poses) - 1):
            p1 = self.current_path.poses[i].pose.position
            p2 = self.current_path.poses[i + 1].pose.position
            segment_dist = self._distance(p1, p2)
            current_dist += segment_dist

            if current_dist >= self.lookahead_distance:
                # Interpolate between p1 and p2
                overshoot = current_dist - self.lookahead_distance
                if segment_dist > 0:
                    alpha = 1.0 - (overshoot / segment_dist)
                    from geometry_msgs.msg import Point
                    lookahead = Point()
                    lookahead.x = p1.x + alpha * (p2.x - p1.x)
                    lookahead.y = p1.y + alpha * (p2.y - p1.y)
                    lookahead.z = 0.0
                    return lookahead

        # If we get here, return the last point
        return self.current_path.poses[-1].pose.position

    def _calculate_steering_command(self, lookahead_point):
        """Calculate steering angle using pure pursuit geometry.

        Args:
            lookahead_point: Target lookahead point on the path

        Returns:
            Tuple of (linear_velocity, angular_velocity)
        """
        if lookahead_point is None or self.current_pose is None:
            return 0.0, 0.0

        # Vector from robot to lookahead point
        dx = lookahead_point.x - self.current_pose.position.x
        dy = lookahead_point.y - self.current_pose.position.y
        distance_to_lookahead = math.sqrt(dx * dx + dy * dy)

        # Check if goal is reached
        if distance_to_lookahead < self.goal_tolerance:
            self.goal_reached = True
            return 0.0, 0.0

        # Robot heading
        robot_yaw = self._quaternion_to_yaw(self.current_pose.orientation)

        # Angle to lookahead point
        angle_to_lookahead = math.atan2(dy, dx)

        # Heading error
        heading_error = angle_to_lookahead - robot_yaw

        # Normalize angle to [-pi, pi]
        while heading_error > math.pi:
            heading_error -= 2 * math.pi
        while heading_error < -math.pi:
            heading_error += 2 * math.pi

        # Pure pursuit curvature: k = 2 * sin(alpha) / Ld
        # where alpha is the heading error and Ld is lookahead distance
        if self.lookahead_distance > 0:
            curvature = (2.0 * math.sin(heading_error)) / self.lookahead_distance
        else:
            curvature = 0.0

        # Linear velocity (reduce when heading error is large)
        if self.use_heading_correction:
            heading_error_abs = abs(heading_error)
            linear_velocity = self.max_linear_velocity * math.cos(heading_error)
            linear_velocity = max(0.0, linear_velocity)  # Don't go backwards
        else:
            linear_velocity = self.max_linear_velocity

        # Angular velocity from curvature: omega = v * k
        angular_velocity = linear_velocity * curvature

        # Clamp velocities
        linear_velocity = min(linear_velocity, self.max_linear_velocity)
        angular_velocity = max(-self.max_angular_velocity,
                               min(angular_velocity, self.max_angular_velocity))

        # Apply scaling
        linear_velocity *= self.linear_velocity_scale
        angular_velocity *= self.angular_velocity_scale

        return linear_velocity, angular_velocity

    def _control_loop(self) -> None:
        """Main control loop executed at regular intervals."""
        if not self.enable_path_following or self.current_pose is None:
            # Publish zero velocity
            cmd = Twist()
            self.pub_cmd_vel.publish(cmd)
            return

        # Find lookahead point
        lookahead_point = self._find_lookahead_point()

        if lookahead_point is None:
            # No valid path, publish zero velocity
            cmd = Twist()
            self.pub_cmd_vel.publish(cmd)
            return

        # Calculate steering command
        linear_vel, angular_vel = self._calculate_steering_command(lookahead_point)

        # Publish velocity command
        cmd = Twist()
        cmd.linear.x = linear_vel
        cmd.angular.z = angular_vel
        self.pub_cmd_vel.publish(cmd)

        # Publish tracking error
        if self.current_path and len(self.current_path.poses) > 0:
            closest_idx, tracking_error, _ = self._find_closest_point_on_path()
            error_msg = Float32(data=tracking_error)
            self.pub_tracking_error.publish(error_msg)


def main(args=None):
    rclpy.init(args=args)
    node = PurePursuitNode()
    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    finally:
        node.destroy_node()
        rclpy.shutdown()


if __name__ == '__main__':
    main()