protoflight 0.1.3

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

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

use crate::{
    autopilot::Autopilot,
    flight::RxMessage,
    tasks::{
        gyro_pid::{GyroPidReceiver, gyro_pid_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};

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>;

#[allow(clippy::expect_used)]
pub fn autopilot_receiver() -> AutopilotReceiver {
    AUTOPILOT_WATCH.receiver().expect("autopilot_receiver failed")
}

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

impl AutopilotContext {
    #[rustfmt::skip]
    pub fn new() -> Self {
        Self {
            gyro_pid_receiver :gyro_pid_receiver(),
            rx_receiver:rx_message_receiver(),
            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(),
        }
    }
}

pub fn init() -> &'static mut AutopilotContext {
    AUTOPILOT_CTX.init(AutopilotContext::new())
}

/// Autopilot Placeholder.
#[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;

        // 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() {
                let vertical_acceleration = gyro_pid_message.acc.z;

                ctx.autopilot.altitude_kalman_filter.predict(vertical_acceleration, delta_t);

                // 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 altitude_hold {
                    let altitude = ctx.autopilot.altitude_kalman_filter.pos();
                    let vertical_speed = ctx.autopilot.altitude_kalman_filter.vel();
                    let throttle_stick = ctx.autopilot.altitude_controller.update(
                        altitude,
                        vertical_speed,
                        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.altitude_kalman_filter.correct_altitude_using_barometer(barometer_message.altitude_m);
            #[cfg(feature = "gps")]
            ctx.autopilot.position_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
        {
            let rangefinder_base_altitude_m = 0.0_f32;
            let altitude = rangefinder_base_altitude_m + rangefinder_message.distance_m;
            ctx.autopilot.altitude_kalman_filter.correct_altitude_using_rangefinder(altitude);
            #[cfg(feature = "gps")]
            ctx.autopilot.position_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.altitude_kalman_filter.correct_altitude_using_gps(gps_position.z);
                } else if position_hold {
                    ctx.autopilot.position_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.
            }
        }

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