mecha10-behavior-runtime 0.9.0

//! Planning node for the "go-to-location" MVP (LAB-4442): drives the robot to a relative
//! navigation goal using odometry-based dead reckoning. No SLAM/map is involved - the goal
//! received on `navigation::GOAL` is expressed in the robot's own frame at the moment the goal
//! was issued (see `SetGoToLocationTargetNode`), and this node tracks progress toward it purely
//! by integrating odometry deltas since that moment.
//!
//! This is the behaviour-tree counterpart to `SetGoToLocationTargetNode`: that node turns a
//! dashboard click into a goal, this node turns a goal into motion. Both are registered
//! independently in `register_builtin_actions` and are meant to run side by side (e.g. under a
//! `Parallel`/`Selector`), the same way `SetInstanceTrackTargetNode`/`InstanceTrackTaskNode` do.

use crate::control::{PidController, PidGains};
use crate::{BehaviorNode, Context, NodeStatus};
use async_trait::async_trait;
use mecha10_core::messages::navigation::NavigationState;
use mecha10_core::messages::{NavigationGoal, NavigationStatus, Odometry, Twist};
use mecha10_core::topics::{actuator, navigation, sensor};
use mecha10_core::Receiver;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::time::Instant;

fn default_position_tolerance_m() -> f32 {
    0.15
}

fn default_max_linear_velocity() -> f32 {
    0.4
}

fn default_max_angular_velocity() -> f32 {
    1.0
}

fn default_heading_lock_threshold_rad() -> f32 {
    // While the heading error is larger than this, spin in place rather than drive forward.
    0.6
}

fn default_linear_gains() -> PidGains {
    PidGains {
        kp: 0.8,
        ki: 0.0,
        kd: 0.0,
    }
}

fn default_angular_gains() -> PidGains {
    PidGains {
        kp: 1.5,
        ki: 0.0,
        kd: 0.1,
    }
}

#[derive(Debug, Serialize, Deserialize)]
struct GoToLocationConfig {
    #[serde(default = "default_position_tolerance_m")]
    position_tolerance_m: f32,
    #[serde(default = "default_max_linear_velocity")]
    max_linear_velocity: f32,
    #[serde(default = "default_max_angular_velocity")]
    max_angular_velocity: f32,
    #[serde(default = "default_heading_lock_threshold_rad")]
    heading_lock_threshold_rad: f32,
    #[serde(default = "default_linear_gains")]
    linear_gains: PidGains,
    #[serde(default = "default_angular_gains")]
    angular_gains: PidGains,
}

impl Default for GoToLocationConfig {
    fn default() -> Self {
        Self {
            position_tolerance_m: default_position_tolerance_m(),
            max_linear_velocity: default_max_linear_velocity(),
            max_angular_velocity: default_max_angular_velocity(),
            heading_lock_threshold_rad: default_heading_lock_threshold_rad(),
            linear_gains: default_linear_gains(),
            angular_gains: default_angular_gains(),
        }
    }
}

/// The odometry pose captured when a goal is accepted - the frame the goal's relative target is
/// expressed in. Everything downstream (`remaining_in_anchor_frame`) is relative to this.
#[derive(Debug, Clone, Copy)]
pub(crate) struct Anchor {
    position: [f32; 2],
    yaw: f32,
}

/// Extracts yaw (rotation about the vertical axis) from a `[x, y, z, w]` quaternion.
pub(crate) fn yaw_from_quaternion(q: [f32; 4]) -> f32 {
    let (x, y, z, w) = (q[0], q[1], q[2], q[3]);
    (2.0 * (w * z + x * y)).atan2(1.0 - 2.0 * (y * y + z * z))
}

/// Wraps an angle to `(-pi, pi]`.
pub(crate) fn normalize_angle(angle: f32) -> f32 {
    let two_pi = std::f32::consts::TAU;
    let mut wrapped = angle % two_pi;
    if wrapped > std::f32::consts::PI {
        wrapped -= two_pi;
    } else if wrapped <= -std::f32::consts::PI {
        wrapped += two_pi;
    }
    wrapped
}

/// Transforms a world-frame point into the `anchor`'s frame (the frame the goal's `target.
/// position` is expressed in) - i.e. "how far has the robot moved, expressed in the direction
/// it was originally facing".
pub(crate) fn world_to_anchor_frame(anchor: &Anchor, world_point: [f32; 2]) -> [f32; 2] {
    let dx = world_point[0] - anchor.position[0];
    let dy = world_point[1] - anchor.position[1];
    let (sin_a, cos_a) = (-anchor.yaw).sin_cos();
    [dx * cos_a - dy * sin_a, dx * sin_a + dy * cos_a]
}

/// Computes the remaining distance and heading error (goal heading minus current heading, both
/// in the anchor frame, wrapped to `(-pi, pi]`) toward `goal_position` from the robot's current
/// `traveled` offset (already in the anchor frame) and `heading_in_anchor_frame`.
pub(crate) fn compute_control(goal_position: [f32; 2], traveled: [f32; 2], heading_in_anchor_frame: f32) -> (f32, f32) {
    let remaining = [goal_position[0] - traveled[0], goal_position[1] - traveled[1]];
    let distance = (remaining[0] * remaining[0] + remaining[1] * remaining[1]).sqrt();
    let desired_heading = remaining[1].atan2(remaining[0]);
    let heading_error = normalize_angle(desired_heading - heading_in_anchor_frame);
    (distance, heading_error)
}

pub struct GoToLocationNode {
    config: GoToLocationConfig,
    goal_receiver: Option<Receiver<NavigationGoal>>,
    odometry_receiver: Option<Receiver<Odometry>>,
    cancel_receiver: Option<Receiver<bool>>,
    goal: Option<NavigationGoal>,
    anchor: Option<Anchor>,
    start_distance: Option<f32>,
    last_tick: Option<Instant>,
    linear_pid: PidController,
    angular_pid: PidController,
}

impl std::fmt::Debug for GoToLocationNode {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("GoToLocationNode")
            .field("has_goal", &self.goal.is_some())
            .field("subscribed", &self.goal_receiver.is_some())
            .finish()
    }
}

impl GoToLocationNode {
    /// Create a new node from JSON configuration. `null` is normalized to `{}`.
    pub fn from_json(config: Value) -> anyhow::Result<Self> {
        let config = if config.is_null() {
            serde_json::json!({})
        } else {
            config
        };
        let cfg: GoToLocationConfig = serde_json::from_value(config)?;
        Ok(Self::with_config(cfg))
    }

    fn with_config(config: GoToLocationConfig) -> Self {
        let linear_pid = PidController::new(config.linear_gains, (0.0, config.max_linear_velocity as f64));
        let angular_pid = PidController::new(
            config.angular_gains,
            (-config.max_angular_velocity as f64, config.max_angular_velocity as f64),
        );
        Self {
            config,
            goal_receiver: None,
            odometry_receiver: None,
            cancel_receiver: None,
            goal: None,
            anchor: None,
            start_distance: None,
            last_tick: None,
            linear_pid,
            angular_pid,
        }
    }

    fn accept_goal(&mut self, goal: NavigationGoal, odometry: &Odometry) {
        let anchor = Anchor {
            position: [odometry.position[0], odometry.position[1]],
            yaw: yaw_from_quaternion(odometry.orientation),
        };
        let goal_position = [goal.target.position[0], goal.target.position[1]];
        self.start_distance = Some((goal_position[0].powi(2) + goal_position[1].powi(2)).sqrt());
        self.goal = Some(goal);
        self.anchor = Some(anchor);
        self.linear_pid.reset();
        self.angular_pid.reset();
        self.last_tick = None;
    }

    fn clear_goal(&mut self) {
        self.goal = None;
        self.anchor = None;
        self.start_distance = None;
        self.last_tick = None;
    }

    async fn publish_stop(&self, ctx: &Context) -> anyhow::Result<()> {
        ctx.publish_to(actuator::CMD_VEL_NAV, &Twist::stop()).await?;
        Ok(())
    }

    async fn publish_status(
        &self,
        ctx: &Context,
        goal_id: &str,
        state: NavigationState,
        distance_remaining: f32,
        progress: f32,
    ) -> anyhow::Result<()> {
        let status = NavigationStatus {
            timestamp: mecha10_core::prelude::now_micros(),
            goal_id: goal_id.to_string(),
            state,
            distance_remaining,
            time_remaining: None,
            progress,
        };
        ctx.publish_to(navigation::STATUS, &status).await?;
        Ok(())
    }

    fn progress(&self, distance_remaining: f32) -> f32 {
        match self.start_distance {
            Some(start) if start > f32::EPSILON => (1.0 - distance_remaining / start).clamp(0.0, 1.0),
            _ => 1.0,
        }
    }
}

#[async_trait]
impl BehaviorNode for GoToLocationNode {
    async fn tick(&mut self, ctx: &Context) -> anyhow::Result<NodeStatus> {
        if self.goal_receiver.is_none() {
            self.goal_receiver = Some(ctx.subscribe(navigation::GOAL).await?);
        }
        if self.odometry_receiver.is_none() {
            self.odometry_receiver = Some(ctx.subscribe(sensor::ODOMETRY).await?);
        }
        if self.cancel_receiver.is_none() {
            self.cancel_receiver = Some(ctx.subscribe(navigation::CANCEL).await?);
        }

        let mut cancelled = false;
        while let Ok(cancel) = self.cancel_receiver.as_mut().expect("just subscribed above").try_recv() {
            cancelled = cancelled || cancel;
        }

        let mut latest_odometry = None;
        while let Ok(odometry) = self
            .odometry_receiver
            .as_mut()
            .expect("just subscribed above")
            .try_recv()
        {
            latest_odometry = Some(odometry);
        }

        let mut latest_goal = None;
        while let Ok(goal) = self.goal_receiver.as_mut().expect("just subscribed above").try_recv() {
            latest_goal = Some(goal);
        }

        if cancelled {
            if let Some(goal) = self.goal.take() {
                self.publish_stop(ctx).await?;
                self.publish_status(ctx, &goal.goal_id, NavigationState::Failed, 0.0, self.progress(0.0))
                    .await?;
            }
            self.clear_goal();
            return Ok(NodeStatus::Failure);
        }

        if let Some(goal) = latest_goal {
            if let Some(odometry) = latest_odometry.as_ref() {
                self.accept_goal(goal, odometry);
            } else {
                tracing::warn!("GoToLocationNode: dropping goal received before any odometry - can't anchor it");
            }
        }

        let Some(goal) = self.goal.clone() else {
            return Ok(NodeStatus::Failure);
        };

        let Some(odometry) = latest_odometry else {
            return Ok(NodeStatus::Running);
        };

        let anchor = self.anchor.expect("a goal is only ever set alongside its anchor");
        let traveled = world_to_anchor_frame(&anchor, [odometry.position[0], odometry.position[1]]);
        let heading_in_anchor_frame = normalize_angle(yaw_from_quaternion(odometry.orientation) - anchor.yaw);
        let goal_position = [goal.target.position[0], goal.target.position[1]];
        let (distance_remaining, heading_error) = compute_control(goal_position, traveled, heading_in_anchor_frame);

        let tolerance = goal.position_tolerance.unwrap_or(self.config.position_tolerance_m);
        if distance_remaining <= tolerance {
            self.publish_stop(ctx).await?;
            self.publish_status(ctx, &goal.goal_id, NavigationState::Succeeded, distance_remaining, 1.0)
                .await?;
            self.clear_goal();
            return Ok(NodeStatus::Success);
        }

        let now = Instant::now();
        let dt = self
            .last_tick
            .map(|last| now.duration_since(last).as_secs_f64())
            .unwrap_or(0.0);
        self.last_tick = Some(now);

        let angular = self.angular_pid.update(heading_error as f64, 0.0, dt) as f32;
        let linear = if heading_error.abs() <= self.config.heading_lock_threshold_rad {
            self.linear_pid.update(distance_remaining as f64, 0.0, dt) as f32
        } else {
            0.0
        };

        ctx.publish_to(actuator::CMD_VEL_NAV, &Twist::new_2d(linear, angular))
            .await?;
        self.publish_status(
            ctx,
            &goal.goal_id,
            NavigationState::Moving,
            distance_remaining,
            self.progress(distance_remaining),
        )
        .await?;

        Ok(NodeStatus::Running)
    }

    fn reset_self(&mut self) {
        self.clear_goal();
        self.linear_pid.reset();
        self.angular_pid.reset();
    }

    fn name(&self) -> &str {
        "GoToLocationNode"
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use mecha10_core::messages::Pose;
    use std::f32::consts::{FRAC_PI_2, PI};

    fn pose_2d(x: f32, y: f32, yaw: f32) -> Pose {
        Pose::new_2d(x, y, yaw)
    }

    #[test]
    fn yaw_from_quaternion_recovers_a_known_yaw() {
        let pose = pose_2d(0.0, 0.0, FRAC_PI_2);
        let yaw = yaw_from_quaternion(pose.orientation);
        assert!((yaw - FRAC_PI_2).abs() < 1e-5);
    }

    #[test]
    fn yaw_from_quaternion_is_zero_for_the_identity_quaternion() {
        assert_eq!(yaw_from_quaternion([0.0, 0.0, 0.0, 1.0]), 0.0);
    }

    #[test]
    fn normalize_angle_wraps_values_outside_pi() {
        assert!((normalize_angle(PI * 1.5) - (-PI * 0.5)).abs() < 1e-5);
        assert!((normalize_angle(-PI * 1.5) - (PI * 0.5)).abs() < 1e-5);
        assert_eq!(normalize_angle(0.0), 0.0);
    }

    #[test]
    fn world_to_anchor_frame_is_identity_when_anchor_is_at_the_origin_facing_forward() {
        let anchor = Anchor {
            position: [0.0, 0.0],
            yaw: 0.0,
        };
        assert_eq!(world_to_anchor_frame(&anchor, [1.0, 2.0]), [1.0, 2.0]);
    }

    #[test]
    fn world_to_anchor_frame_rotates_by_the_negative_anchor_yaw() {
        // Anchor facing +90deg (yaw = pi/2): the world's +y axis becomes the anchor's +x axis.
        let anchor = Anchor {
            position: [0.0, 0.0],
            yaw: FRAC_PI_2,
        };
        let transformed = world_to_anchor_frame(&anchor, [0.0, 1.0]);
        assert!((transformed[0] - 1.0).abs() < 1e-5);
        assert!(transformed[1].abs() < 1e-5);
    }

    #[test]
    fn world_to_anchor_frame_subtracts_the_anchor_position_first() {
        let anchor = Anchor {
            position: [5.0, 5.0],
            yaw: 0.0,
        };
        assert_eq!(world_to_anchor_frame(&anchor, [6.0, 5.0]), [1.0, 0.0]);
    }

    #[test]
    fn compute_control_reports_zero_distance_and_error_once_at_the_goal() {
        let (distance, heading_error) = compute_control([2.0, 0.0], [2.0, 0.0], 0.0);
        assert_eq!(distance, 0.0);
        assert_eq!(heading_error, 0.0);
    }

    #[test]
    fn compute_control_points_straight_ahead_when_the_goal_is_directly_ahead() {
        let (distance, heading_error) = compute_control([3.0, 0.0], [0.0, 0.0], 0.0);
        assert_eq!(distance, 3.0);
        assert_eq!(heading_error, 0.0);
    }

    #[test]
    fn compute_control_reports_a_positive_heading_error_for_a_goal_to_the_left() {
        let (_, heading_error) = compute_control([0.0, 3.0], [0.0, 0.0], 0.0);
        assert!((heading_error - FRAC_PI_2).abs() < 1e-5);
    }

    #[test]
    fn compute_control_accounts_for_the_robots_current_heading() {
        // Goal is straight ahead in the anchor frame, but the robot has already turned 90deg
        // left - so from its current heading, the goal is now 90deg to the right.
        let (_, heading_error) = compute_control([3.0, 0.0], [0.0, 0.0], FRAC_PI_2);
        assert!((heading_error + FRAC_PI_2).abs() < 1e-5);
    }

    #[test]
    fn from_json_defaults_all_fields() {
        let node = GoToLocationNode::from_json(serde_json::json!({})).unwrap();
        assert_eq!(node.config.position_tolerance_m, 0.15);
        assert_eq!(node.config.max_linear_velocity, 0.4);
        assert_eq!(node.config.max_angular_velocity, 1.0);
    }

    #[test]
    fn from_json_accepts_null_the_same_as_an_empty_object() {
        let node = GoToLocationNode::from_json(Value::Null).unwrap();
        assert_eq!(node.config.position_tolerance_m, 0.15);
    }

    #[test]
    fn from_json_reads_overridden_fields() {
        let node = GoToLocationNode::from_json(serde_json::json!({
            "position_tolerance_m": 0.5,
            "max_linear_velocity": 1.0,
            "max_angular_velocity": 2.0
        }))
        .unwrap();
        assert_eq!(node.config.position_tolerance_m, 0.5);
        assert_eq!(node.config.max_linear_velocity, 1.0);
        assert_eq!(node.config.max_angular_velocity, 2.0);
    }

    fn sample_goal(x: f32, y: f32) -> NavigationGoal {
        NavigationGoal::new_2d("test-goal".to_string(), x, y, 0.0)
    }

    fn sample_odometry(x: f32, y: f32, yaw: f32) -> Odometry {
        let pose = pose_2d(x, y, yaw);
        Odometry {
            timestamp: 0,
            position: pose.position,
            orientation: pose.orientation,
            linear_velocity: [0.0, 0.0, 0.0],
            angular_velocity: [0.0, 0.0, 0.0],
            frame_id: None,
            child_frame_id: None,
        }
    }

    #[test]
    fn accept_goal_anchors_on_the_odometry_pose_at_acceptance_time() {
        let mut node = GoToLocationNode::from_json(Value::Null).unwrap();
        let goal = sample_goal(2.0, 0.0);
        let odometry = sample_odometry(1.0, 1.0, FRAC_PI_2);

        node.accept_goal(goal, &odometry);

        let anchor = node.anchor.expect("anchor set");
        assert_eq!(anchor.position, [1.0, 1.0]);
        assert!((anchor.yaw - FRAC_PI_2).abs() < 1e-5);
        assert_eq!(node.start_distance, Some(2.0));
    }

    #[test]
    fn clear_goal_resets_tracking_state() {
        let mut node = GoToLocationNode::from_json(Value::Null).unwrap();
        node.accept_goal(sample_goal(1.0, 0.0), &sample_odometry(0.0, 0.0, 0.0));
        assert!(node.goal.is_some());

        node.clear_goal();

        assert!(node.goal.is_none());
        assert!(node.anchor.is_none());
        assert!(node.start_distance.is_none());
    }

    #[test]
    fn progress_tracks_fraction_of_distance_covered() {
        let mut node = GoToLocationNode::from_json(Value::Null).unwrap();
        node.accept_goal(sample_goal(4.0, 0.0), &sample_odometry(0.0, 0.0, 0.0));

        assert_eq!(node.progress(4.0), 0.0);
        assert_eq!(node.progress(2.0), 0.5);
        assert_eq!(node.progress(0.0), 1.0);
    }

    #[test]
    fn progress_is_complete_when_there_is_no_tracked_start_distance() {
        let node = GoToLocationNode::from_json(Value::Null).unwrap();
        assert_eq!(node.progress(5.0), 1.0);
    }
}