protoflight 0.1.5

#![cfg(feature = "autopilot")]

use vqm::Vector3f32;

use embassy_sync::{
    blocking_mutex::raw::CriticalSectionRawMutex,
    signal::Signal,
    watch::{Receiver, Sender, Watch},
};
use static_cell::StaticCell;

use crate::{
    autopilot::Autopilot,
    flight::RxMessage,
    tasks::{
        errors::TaskInitError,
        failsafe::{FailsafeSubscriber, failsafe_subscriber},
        gyro_pid::{GyroPidReceiver, ImuIntegralReceiver, gyro_pid_receiver, imu_integral_receiver},
        rx::{RxMessageReceiver, rx_message_receiver},
    },
};

static AUTOPILOT_CTX: StaticCell<AutopilotContext> = StaticCell::new();

#[cfg(any(feature = "barometer", feature = "gps", feature = "optical_flow", feature = "rangefinder"))]
use {crate::flight::RcControls, radio_controllers::RcMode};

#[cfg(feature = "barometer")]
use crate::tasks::barometer::{BarometerSubscriber, barometer_subscriber};

#[cfg(feature = "gps")]
use crate::{
    gps::GpsMessage,
    tasks::gps::{GpsSubscriber, gps_subscriber},
};

#[cfg(feature = "optical_flow")]
use crate::tasks::optical_flow::{OpticalFlowSubscriber, optical_flow_subscriber};

#[cfg(feature = "rangefinder")]
use crate::tasks::rangefinder::{RangefinderSubscriber, rangefinder_subscriber};

pub static KALMAN_ACC_BIAS_SIGNAL: Signal<CriticalSectionRawMutex, KalmanAccBiasNedMessage> = Signal::new();

const AUTOPILOT_WATCH_COUNT: usize = 1;
static AUTOPILOT_WATCH: Watch<CriticalSectionRawMutex, RxMessage, AUTOPILOT_WATCH_COUNT> = Watch::new();

type AutopilotSender = Sender<'static, CriticalSectionRawMutex, RxMessage, AUTOPILOT_WATCH_COUNT>;
pub fn autopilot_sender() -> AutopilotSender {
    AUTOPILOT_WATCH.sender()
}

pub type AutopilotReceiver = Receiver<'static, CriticalSectionRawMutex, RxMessage, AUTOPILOT_WATCH_COUNT>;

pub fn autopilot_receiver() -> Result<AutopilotReceiver, TaskInitError> {
    AUTOPILOT_WATCH.receiver().ok_or(TaskInitError::AutopilotWatchCountTooLow)
}

#[derive(Clone, Copy, Debug, PartialEq)]
#[cfg_attr(feature = "std", derive(derive_more::Display))]
pub struct KalmanAccBiasNedMessage {
    pub acc_bias_ned_mps2: Vector3f32,
}

impl Default for KalmanAccBiasNedMessage {
    fn default() -> Self {
        Self::new()
    }
}

impl KalmanAccBiasNedMessage {
    pub const fn new() -> Self {
        Self { acc_bias_ned_mps2: Vector3f32 { x: 0.0, y: 0.0, z: 0.0 } }
    }
}

/// Context for Autopilot task.
#[rustfmt::skip]
pub struct AutopilotContext {
    pub gyro_pid_receiver: GyroPidReceiver,
    pub imu_integral_receiver: ImuIntegralReceiver,
    pub rx_receiver: RxMessageReceiver,
    pub failsafe_subscriber: FailsafeSubscriber,
    pub autopilot_sender: AutopilotSender,
    pub autopilot: Autopilot,
    #[cfg(feature = "barometer")] pub barometer_subscriber: BarometerSubscriber,
    #[cfg(feature = "gps")] pub gps_subscriber: GpsSubscriber,
    #[cfg(feature = "optical_flow")] pub optical_flow_subscriber: OpticalFlowSubscriber,
    #[cfg(feature = "rangefinder")] pub rangefinder_subscriber: RangefinderSubscriber,
}

pub fn init() -> Result<&'static mut AutopilotContext, TaskInitError> {
    #[rustfmt::skip]
    let ctx = AutopilotContext {
        gyro_pid_receiver: gyro_pid_receiver()?,
        imu_integral_receiver: imu_integral_receiver()?,
        rx_receiver: rx_message_receiver()?,
        failsafe_subscriber: failsafe_subscriber()?,
        autopilot_sender: autopilot_sender(),
        autopilot: Autopilot::new(),
        #[cfg(feature = "barometer")] barometer_subscriber: barometer_subscriber()?,
        #[cfg(feature = "gps")] gps_subscriber: gps_subscriber()?,
        #[cfg(feature = "optical_flow")] optical_flow_subscriber: optical_flow_subscriber()?,
        #[cfg(feature = "rangefinder")] rangefinder_subscriber: rangefinder_subscriber()?,
    };
    Ok(AUTOPILOT_CTX.init(ctx))
}

// TODO: make this run on receipt of the integrated_imu message, rather than on a timer.
/// Autopilot Placeholder.
#[allow(clippy::too_many_lines)]
#[allow(unused)]
#[embassy_executor::task]
pub async fn run(ctx: &'static mut AutopilotContext) {
    let mut ticker = embassy_time::Ticker::every(embassy_time::Duration::from_millis(1));
    let delta_t = 0.001;
    let mut loop_count: u32 = 0;

    #[allow(unused)]
    let mut altitude_hold = false;
    #[cfg(feature = "gps")]
    let mut position_hold = false;

    log::info!("   AUTOPILOT: task started");
    loop {
        ticker.next().await;

        if let Some(imu_integral_message) = ctx.imu_integral_receiver.try_get() {
            let imu_integral = imu_integral_message.imu_integral;
            // call the Kalman filter to predict using delta_vel and delta_pos from the gyro/pid task.
            #[cfg(not(any(feature = "gps", feature = "optical_flow")))]
            {
                ctx.autopilot.kalman_filter.predict(
                    imu_integral.delta_vel_ned_mps.z,
                    imu_integral.delta_pos_ned_m.z,
                    delta_t,
                );
                // signal the acc bias so the gyro_pid task use it to update the IMU integrator.
                let acc_bias_ned_mps2 = Vector3f32 { x: 0.0, y: 0.0, z: ctx.autopilot.kalman_filter.acc_bias_mps2() };
                KALMAN_ACC_BIAS_SIGNAL.signal(KalmanAccBiasNedMessage { acc_bias_ned_mps2 });
            }
            #[cfg(any(feature = "gps", feature = "optical_flow"))]
            {
                ctx.autopilot.kalman_filter.predict(
                    imu_integral.delta_vel_ned_mps,
                    imu_integral.delta_pos_ned_m,
                    delta_t,
                );
                // signal the acc bias so the gyro_pid task use it to update the IMU integrator.
                let acc_bias_ned_mps2 = ctx.autopilot.kalman_filter.acc_bias_mps2();
                KALMAN_ACC_BIAS_SIGNAL.signal(KalmanAccBiasNedMessage { acc_bias_ned_mps2 });
            }
        }

        // for the autopilot to provide any functionality, it has to have at least one of: barometer, gps, optical_flow, or rangefinder.
        #[cfg(any(feature = "barometer", feature = "gps", feature = "optical_flow", feature = "rangefinder"))]
        {
            if let Some(gyro_pid_message) = ctx.gyro_pid_receiver.try_get() {
                use embassy_sync::pubsub::WaitResult;

                // Check if the rc_modes have changed.
                if let Some(rx_message) = ctx.rx_receiver.try_changed() {
                    let rc_modes = rx_message.rc_modes;
                    altitude_hold = rc_modes.test(RcMode::ALTITUDE_HOLD);
                    #[cfg(feature = "gps")]
                    {
                        position_hold = rc_modes.test(RcMode::POSITION_HOLD)
                            | rc_modes.test(RcMode::GPS_RESCUE)
                            | rc_modes.test(RcMode::AUTOPILOT);
                    }
                }
                if let Some(WaitResult::Message(failsafe_message)) = ctx.failsafe_subscriber.try_next_message() {
                    _ = failsafe_message;
                }
                if altitude_hold {
                    #[cfg(not(any(feature = "gps", feature = "optical_flow")))]
                    let (altitude_m, vertical_speed_mps) =
                        (ctx.autopilot.kalman_filter.pos_m(), ctx.autopilot.kalman_filter.vel_mps());
                    #[cfg(any(feature = "gps", feature = "optical_flow"))]
                    let (altitude_m, vertical_speed_mps) =
                        (ctx.autopilot.kalman_filter.pos_m().z, ctx.autopilot.kalman_filter.vel_mps().z);

                    let throttle_stick = ctx.autopilot.altitude_controller.update(
                        altitude_m,
                        vertical_speed_mps,
                        gyro_pid_message.orientation,
                        delta_t,
                    );

                    // Send the flight control message. This will be picked by the radio task.
                    let rx_message = RxMessage {
                        rc_controls: RcControls { throttle_stick, ..Default::default() },
                        ..Default::default()
                    };
                    ctx.autopilot_sender.send(rx_message);
                }
            }
        }

        // Note: `try_next_message` is the non-blocking polling form.
        // If there is a message, it is removed and processed immediately.
        // Lagged (missed) messages are ignored.
        #[cfg(feature = "barometer")]
        if let Some(wait_result) = ctx.barometer_subscriber.try_next_message()
            && let embassy_sync::pubsub::WaitResult::Message(barometer_message) = wait_result
        {
            _ = ctx.autopilot.kalman_filter.correct_altitude_using_barometer(barometer_message.altitude_m);
        }

        #[cfg(feature = "rangefinder")]
        if let Some(wait_result) = ctx.rangefinder_subscriber.try_next_message()
            && let embassy_sync::pubsub::WaitResult::Message(rangefinder_message) = wait_result
        {
            // TODO: set rangefinder_base_altitude_m
            let rangefinder_base_altitude_m = 0.0_f32;
            let altitude = rangefinder_base_altitude_m + rangefinder_message.distance_m;
            _ = ctx.autopilot.kalman_filter.correct_altitude_using_rangefinder(altitude);
        }

        #[cfg(feature = "gps")]
        if let Some(wait_result) = ctx.gps_subscriber.try_next_message()
            && let embassy_sync::pubsub::WaitResult::Message(event) = wait_result
        {
            // TODO: choose position or altitude kalman filter based on settings
            if let GpsMessage::Position(gps_position) = event {
                if altitude_hold {
                    _ = ctx.autopilot.kalman_filter.correct_altitude_using_gps(gps_position.z);
                } else if position_hold {
                    _ = ctx.autopilot.kalman_filter.correct_position_using_gps(gps_position);
                }
            } else {
                // Message type of interest to other subscribers, but not to me so intentionally do nothing,
                // this consumes the message and removes it from the queue.
            }
        }

        #[cfg(feature = "optical_flow")]
        if let Some(wait_result) = ctx.optical_flow_subscriber.try_next_message()
            && let embassy_sync::pubsub::WaitResult::Message(optical_flow_message) = wait_result
        {
            _ = optical_flow_message;
        }

        if loop_count.is_multiple_of(200) {
            log::info!("       AUTOPILOT:loop {loop_count}");
        }
        loop_count = loop_count.wrapping_add(1);
    }
}