mecha10-behavior-runtime 0.8.2

//! Flight-stabilization action nodes for rotorcraft: `TakeoffNode` climbs to a target altitude,
//! `LandNode` descends back to the ground, both using a PID altitude-hold loop
//! ([`crate::control::PidController`]) against odometry feedback to compute per-rotor thrust.

use crate::control::{PidController, PidGains};
use crate::{BehaviorNode, Context, NodeStatus};
use async_trait::async_trait;
use mecha10_core::messages::Odometry;
use mecha10_core::topics::Topic;
use mecha10_core::Receiver;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::time::{Duration, Instant};

/// Shared configuration fields for [`TakeoffNode`]/[`LandNode`] - everything except the
/// altitude each one holds as its setpoint.
#[derive(Debug, Serialize, Deserialize)]
struct AltitudeHoldConfig {
    /// Topic to publish per-rotor thrust commands to (a `Vec<f32>`, one entry per rotor, all
    /// equal since this is a pure vertical-thrust hold with no attitude correction).
    #[serde(default = "default_motor_cmd_topic")]
    motor_cmd_topic: String,
    /// Topic to read altitude feedback from (`Odometry.position[2]`).
    #[serde(default = "default_odometry_topic")]
    odometry_topic: String,
    /// Number of rotors to publish a thrust value for - matches the 4-prop drone template.
    #[serde(default = "default_rotor_count")]
    rotor_count: usize,
    /// Baseline thrust (fraction of `max_thrust`) assumed to roughly counteract gravity at
    /// hover; the PID output is added on top of this rather than driving thrust from zero.
    #[serde(default = "default_hover_thrust")]
    hover_thrust: f64,
    #[serde(default = "default_min_thrust")]
    min_thrust: f64,
    #[serde(default = "default_max_thrust")]
    max_thrust: f64,
    /// Altitude error (meters) within which the target is considered reached.
    #[serde(default = "default_tolerance_m")]
    tolerance_m: f64,
    /// Give up and report `Failure` if the target altitude hasn't been reached within this many
    /// seconds (e.g. the drone is stuck, or there's no odometry feed at all).
    #[serde(default = "default_timeout_secs")]
    timeout_secs: f64,
    #[serde(default)]
    pid: PidGains,
}

fn default_motor_cmd_topic() -> String {
    "/actuator/motor_cmd".to_string()
}

fn default_odometry_topic() -> String {
    "/sensor/odometry".to_string()
}

fn default_rotor_count() -> usize {
    4
}

fn default_hover_thrust() -> f64 {
    0.5
}

fn default_min_thrust() -> f64 {
    0.0
}

fn default_max_thrust() -> f64 {
    1.0
}

fn default_tolerance_m() -> f64 {
    0.05
}

fn default_timeout_secs() -> f64 {
    15.0
}

/// Outcome of one [`AltitudeHoldLoop::step`] call, separated from `tick()`'s `ctx` wiring
/// (subscribe/publish) so the altitude/PID decision logic is unit-testable without a live
/// `Context` - same split as `SensorCheckNode::evaluate` vs. `SensorCheckNode::tick`.
#[derive(Debug, PartialEq)]
enum AltitudeStep {
    /// Not reached yet - publish this thrust fraction to every rotor and keep running.
    Hold(f64),
    /// Within tolerance of the target altitude.
    Reached,
    /// `timeout_secs` elapsed without reaching the target altitude.
    TimedOut,
}

/// Shared altitude-hold control loop used by both [`TakeoffNode`] and [`LandNode`].
struct AltitudeHoldLoop {
    motor_cmd_topic: String,
    odometry_topic: String,
    rotor_count: usize,
    hover_thrust: f64,
    min_thrust: f64,
    max_thrust: f64,
    tolerance_m: f64,
    timeout: Duration,
    target_altitude_m: f64,
    pid: PidController,
    receiver: Option<Receiver<Odometry>>,
    current_altitude_m: Option<f64>,
    last_tick: Option<Instant>,
    started_at: Option<Instant>,
}

impl std::fmt::Debug for AltitudeHoldLoop {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("AltitudeHoldLoop")
            .field("motor_cmd_topic", &self.motor_cmd_topic)
            .field("odometry_topic", &self.odometry_topic)
            .field("target_altitude_m", &self.target_altitude_m)
            .field("current_altitude_m", &self.current_altitude_m)
            .field("subscribed", &self.receiver.is_some())
            .finish()
    }
}

impl AltitudeHoldLoop {
    fn new(config: AltitudeHoldConfig, target_altitude_m: f64) -> Self {
        let pid = PidController::new(config.pid, (config.min_thrust, config.max_thrust))
            .with_integral_limits((config.min_thrust, config.max_thrust));

        Self {
            motor_cmd_topic: config.motor_cmd_topic,
            odometry_topic: config.odometry_topic,
            rotor_count: config.rotor_count,
            hover_thrust: config.hover_thrust,
            min_thrust: config.min_thrust,
            max_thrust: config.max_thrust,
            tolerance_m: config.tolerance_m,
            timeout: Duration::from_secs_f64(config.timeout_secs),
            target_altitude_m,
            pid,
            receiver: None,
            current_altitude_m: None,
            last_tick: None,
            started_at: None,
        }
    }

    /// Drain the odometry receiver (keeping only the most recent reading) and advance the PID
    /// loop by one step. Pure with respect to messaging infrastructure - operates on the
    /// already-subscribed `Receiver` and wall-clock time only.
    fn step(&mut self) -> AltitudeStep {
        if let Some(receiver) = self.receiver.as_mut() {
            while let Ok(odom) = receiver.try_recv() {
                self.current_altitude_m = Some(odom.position[2] as f64);
            }
        }

        let now = Instant::now();
        let started_at = *self.started_at.get_or_insert(now);
        if now.duration_since(started_at) >= self.timeout {
            return AltitudeStep::TimedOut;
        }

        let Some(altitude) = self.current_altitude_m else {
            return AltitudeStep::Hold(self.hover_thrust);
        };

        if (self.target_altitude_m - altitude).abs() <= self.tolerance_m {
            return AltitudeStep::Reached;
        }

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

        let correction = self.pid.update(self.target_altitude_m, altitude, dt);
        let thrust = (self.hover_thrust + correction).clamp(self.min_thrust, self.max_thrust);
        AltitudeStep::Hold(thrust)
    }

    async fn publish_thrust(&self, ctx: &Context, thrust: f64) -> anyhow::Result<()> {
        let thrusts = vec![thrust as f32; self.rotor_count];
        let topic_static: &'static str = Box::leak(self.motor_cmd_topic.clone().into_boxed_str());
        ctx.publish_to(Topic::<Vec<f32>>::new(topic_static), &thrusts).await?;
        Ok(())
    }

    async fn ensure_subscribed(&mut self, ctx: &Context) -> anyhow::Result<()> {
        if self.receiver.is_none() {
            self.receiver = Some(ctx.subscribe_raw::<Odometry>(&self.odometry_topic).await?);
        }
        Ok(())
    }

    fn reset(&mut self) {
        self.pid.reset();
        self.current_altitude_m = None;
        self.last_tick = None;
        self.started_at = None;
    }

    #[cfg(test)]
    fn with_receiver(config: AltitudeHoldConfig, target_altitude_m: f64, receiver: Receiver<Odometry>) -> Self {
        let mut loop_ = Self::new(config, target_altitude_m);
        loop_.receiver = Some(receiver);
        loop_
    }
}

/// Climbs from the ground to `target_altitude_m` using a PID altitude-hold loop, publishing
/// per-rotor thrust to `motor_cmd_topic` until the target is reached (`Success`) or
/// `timeout_secs` elapses without reaching it (`Failure`).
pub struct TakeoffNode {
    control: AltitudeHoldLoop,
}

#[derive(Debug, Serialize, Deserialize)]
struct TakeoffConfig {
    target_altitude_m: f64,
    #[serde(flatten)]
    shared: AltitudeHoldConfig,
}

impl std::fmt::Debug for TakeoffNode {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("TakeoffNode").field("control", &self.control).finish()
    }
}

impl TakeoffNode {
    /// Create a new TakeoffNode from JSON configuration
    pub fn from_json(config: Value) -> anyhow::Result<Self> {
        let cfg: TakeoffConfig = serde_json::from_value(config)?;
        Ok(Self {
            control: AltitudeHoldLoop::new(cfg.shared, cfg.target_altitude_m),
        })
    }
}

#[async_trait]
impl BehaviorNode for TakeoffNode {
    async fn tick(&mut self, ctx: &Context) -> anyhow::Result<NodeStatus> {
        self.control.ensure_subscribed(ctx).await?;

        match self.control.step() {
            AltitudeStep::Hold(thrust) => {
                self.control.publish_thrust(ctx, thrust).await?;
                Ok(NodeStatus::Running)
            }
            AltitudeStep::Reached => {
                tracing::debug!("TakeoffNode: reached target altitude");
                Ok(NodeStatus::Success)
            }
            AltitudeStep::TimedOut => {
                tracing::warn!("TakeoffNode: timed out before reaching target altitude");
                Ok(NodeStatus::Failure)
            }
        }
    }

    fn reset_self(&mut self) {
        self.control.reset();
    }

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

/// Descends from the current altitude down to `ground_altitude_m` (default `0.0`) using the
/// same PID altitude-hold loop as [`TakeoffNode`], publishing a zero-thrust command once
/// touchdown is detected so the rotors don't keep holding hover thrust at ground level.
pub struct LandNode {
    control: AltitudeHoldLoop,
}

#[derive(Debug, Serialize, Deserialize)]
struct LandConfig {
    #[serde(default)]
    ground_altitude_m: f64,
    #[serde(flatten)]
    shared: AltitudeHoldConfig,
}

impl std::fmt::Debug for LandNode {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("LandNode").field("control", &self.control).finish()
    }
}

impl LandNode {
    /// Create a new LandNode from JSON configuration
    pub fn from_json(config: Value) -> anyhow::Result<Self> {
        let cfg: LandConfig = serde_json::from_value(config)?;
        Ok(Self {
            control: AltitudeHoldLoop::new(cfg.shared, cfg.ground_altitude_m),
        })
    }
}

#[async_trait]
impl BehaviorNode for LandNode {
    async fn tick(&mut self, ctx: &Context) -> anyhow::Result<NodeStatus> {
        self.control.ensure_subscribed(ctx).await?;

        match self.control.step() {
            AltitudeStep::Hold(thrust) => {
                self.control.publish_thrust(ctx, thrust).await?;
                Ok(NodeStatus::Running)
            }
            AltitudeStep::Reached => {
                tracing::debug!("LandNode: touched down, cutting rotor thrust");
                self.control.publish_thrust(ctx, 0.0).await?;
                Ok(NodeStatus::Success)
            }
            AltitudeStep::TimedOut => {
                tracing::warn!("LandNode: timed out before reaching ground altitude");
                Ok(NodeStatus::Failure)
            }
        }
    }

    fn reset_self(&mut self) {
        self.control.reset();
    }

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

#[cfg(test)]
mod tests {
    use super::*;
    use mecha10_core::Receiver;
    use serde_json::json;
    use tokio::sync::mpsc;

    fn default_config() -> AltitudeHoldConfig {
        AltitudeHoldConfig {
            motor_cmd_topic: default_motor_cmd_topic(),
            odometry_topic: default_odometry_topic(),
            rotor_count: 4,
            hover_thrust: 0.5,
            min_thrust: 0.0,
            max_thrust: 1.0,
            tolerance_m: 0.05,
            timeout_secs: 15.0,
            pid: PidGains {
                kp: 0.5,
                ki: 0.1,
                kd: 0.05,
            },
        }
    }

    fn odometry_at(altitude: f32) -> Odometry {
        Odometry {
            timestamp: 0,
            position: [0.0, 0.0, altitude],
            orientation: [0.0, 0.0, 0.0, 1.0],
            linear_velocity: [0.0, 0.0, 0.0],
            angular_velocity: [0.0, 0.0, 0.0],
            frame_id: None,
            child_frame_id: None,
        }
    }

    fn loop_with_channel(target_altitude_m: f64) -> (AltitudeHoldLoop, mpsc::UnboundedSender<Odometry>) {
        let (tx, rx) = mpsc::unbounded_channel();
        let control = AltitudeHoldLoop::with_receiver(default_config(), target_altitude_m, Receiver::from(rx));
        (control, tx)
    }

    #[test]
    fn no_reading_yet_holds_at_hover_thrust() {
        let (mut control, _tx) = loop_with_channel(1.0);
        assert_eq!(control.step(), AltitudeStep::Hold(0.5));
    }

    #[test]
    fn below_target_climbs_with_thrust_above_hover() {
        let (mut control, tx) = loop_with_channel(1.0);
        tx.send(odometry_at(0.0)).unwrap();

        match control.step() {
            AltitudeStep::Hold(thrust) => assert!(thrust > 0.5, "expected climb thrust above hover, got {thrust}"),
            other => panic!("expected Hold, got {other:?}"),
        }
    }

    #[test]
    fn within_tolerance_reports_reached() {
        let (mut control, tx) = loop_with_channel(1.0);
        tx.send(odometry_at(0.98)).unwrap();
        assert_eq!(control.step(), AltitudeStep::Reached);
    }

    #[test]
    fn exact_target_reports_reached() {
        let (mut control, tx) = loop_with_channel(1.0);
        tx.send(odometry_at(1.0)).unwrap();
        assert_eq!(control.step(), AltitudeStep::Reached);
    }

    #[test]
    fn overshoot_above_target_still_within_tolerance_is_reached() {
        let (mut control, tx) = loop_with_channel(1.0);
        tx.send(odometry_at(1.03)).unwrap();
        assert_eq!(control.step(), AltitudeStep::Reached);
    }

    #[test]
    fn thrust_is_clamped_to_configured_max() {
        let mut config = default_config();
        config.pid.kp = 100.0;
        let (tx, rx) = mpsc::unbounded_channel();
        let mut control = AltitudeHoldLoop::with_receiver(config, 100.0, Receiver::from(rx));
        tx.send(odometry_at(0.0)).unwrap();

        match control.step() {
            AltitudeStep::Hold(thrust) => assert!(thrust <= 1.0),
            other => panic!("expected Hold, got {other:?}"),
        }
    }

    #[test]
    fn later_reading_overrides_earlier_one() {
        let (mut control, tx) = loop_with_channel(1.0);
        tx.send(odometry_at(0.0)).unwrap();
        tx.send(odometry_at(1.0)).unwrap();

        assert_eq!(control.step(), AltitudeStep::Reached);
    }

    #[test]
    fn timed_out_after_timeout_elapses_without_reaching_target() {
        let mut config = default_config();
        config.timeout_secs = 0.0;
        let (tx, rx) = mpsc::unbounded_channel();
        let mut control = AltitudeHoldLoop::with_receiver(config, 1.0, Receiver::from(rx));
        tx.send(odometry_at(0.0)).unwrap();

        assert_eq!(control.step(), AltitudeStep::TimedOut);
    }

    #[test]
    fn reset_clears_altitude_and_pid_state() {
        let (mut control, tx) = loop_with_channel(1.0);
        tx.send(odometry_at(1.0)).unwrap();
        assert_eq!(control.step(), AltitudeStep::Reached);

        control.reset();
        assert_eq!(control.current_altitude_m, None);
    }

    #[test]
    fn takeoff_from_json_parses_defaults() {
        let config = json!({ "target_altitude_m": 1.5 });
        assert!(TakeoffNode::from_json(config).is_ok());
    }

    #[test]
    fn takeoff_from_json_requires_target_altitude() {
        let config = json!({});
        assert!(TakeoffNode::from_json(config).is_err());
    }

    #[test]
    fn land_from_json_defaults_ground_altitude_to_zero() {
        let config = json!({});
        let node = LandNode::from_json(config).unwrap();
        assert_eq!(node.control.target_altitude_m, 0.0);
    }

    #[test]
    fn land_from_json_accepts_custom_ground_altitude() {
        let config = json!({ "ground_altitude_m": 0.2 });
        let node = LandNode::from_json(config).unwrap();
        assert_eq!(node.control.target_altitude_m, 0.2);
    }
}