protoflight 0.1.5

use embassy_sync::{
    blocking_mutex::raw::CriticalSectionRawMutex,
    pubsub::WaitResult,
    watch::{Receiver, Sender, Watch},
};

use static_cell::StaticCell;

use imu_sensors::{AccFullScale, AccUnits, GyroFullScale, GyroUnits, ImuDevice, ImuError};
use motor_mixers::MotorMixerMessage;
#[cfg(feature = "autopilot")]
use sensor_fusion::ImuIntegrator;
use sensor_fusion::{MadgwickFilterf32, SensorFusion};
use simple_bitset::BitSet64;

use crate::{
    boards::targets::BoardImu,
    config::{FastConfigSubscriber, fast_config_subscriber},
    flight::{FilterAccGyro, FlightController, ImuFilterBank, ImuFilterBankConfig, RcControls, VehicleControl},
    tasks::{
        GyroPidMessage, ImuIntegralMessage, SetpointMessage,
        errors::TaskInitError,
        motor_mixer::MOTOR_MIXER_SIGNAL,
        rx::{RxMessageReceiver, rx_message_receiver},
    },
};

#[cfg(feature = "gps")]
use crate::tasks::gps::GPS_YAW_HEADING_SIGNAL;

// The gyro_pid watch has three clients: the blackbox, the autopilot, and the OSD.
const GYRO_PID_WATCH_COUNT: usize = 3;
// Watch<Mutex, DataType, MaxReceivers>
static GYRO_PID_WATCH: Watch<CriticalSectionRawMutex, GyroPidMessage, GYRO_PID_WATCH_COUNT> = Watch::new();

// Type aliases make the function signatures much easier to read.
type GyroPidSender = Sender<'static, CriticalSectionRawMutex, GyroPidMessage, GYRO_PID_WATCH_COUNT>;
pub fn gyro_pid_sender() -> GyroPidSender {
    GYRO_PID_WATCH.sender()
}

#[allow(unused)]
pub type GyroPidReceiver = Receiver<'static, CriticalSectionRawMutex, GyroPidMessage, GYRO_PID_WATCH_COUNT>;
#[allow(unused)]
pub fn gyro_pid_receiver() -> Result<GyroPidReceiver, TaskInitError> {
    GYRO_PID_WATCH.receiver().ok_or(TaskInitError::GyroPidWatchCountTooLow)
}

const IMU_INTEGRAL_WATCH_COUNT: usize = 1;
static IMU_INTEGRAL_WATCH: Watch<CriticalSectionRawMutex, ImuIntegralMessage, IMU_INTEGRAL_WATCH_COUNT> = Watch::new();

type ImuIntegralSender = Sender<'static, CriticalSectionRawMutex, ImuIntegralMessage, IMU_INTEGRAL_WATCH_COUNT>;
#[allow(unused)]
pub fn imu_integral_sender() -> ImuIntegralSender {
    IMU_INTEGRAL_WATCH.sender()
}

pub type ImuIntegralReceiver = Receiver<'static, CriticalSectionRawMutex, ImuIntegralMessage, IMU_INTEGRAL_WATCH_COUNT>;
#[allow(unused)]
pub fn imu_integral_receiver() -> Result<ImuIntegralReceiver, TaskInitError> {
    IMU_INTEGRAL_WATCH.receiver().ok_or(TaskInitError::ImuIntegralWatchCountTooLow)
}

const SETPOINT_WATCH_COUNT: usize = 3;
static SETPOINT_WATCH: Watch<CriticalSectionRawMutex, SetpointMessage, SETPOINT_WATCH_COUNT> = Watch::new();

type SetpointSender = Sender<'static, CriticalSectionRawMutex, SetpointMessage, SETPOINT_WATCH_COUNT>;
pub fn setpoint_sender() -> SetpointSender {
    SETPOINT_WATCH.sender()
}

#[allow(unused)]
pub type SetpointReceiver = Receiver<'static, CriticalSectionRawMutex, SetpointMessage, SETPOINT_WATCH_COUNT>;
#[allow(unused)]
pub fn setpoint_receiver() -> Result<SetpointReceiver, TaskInitError> {
    SETPOINT_WATCH.receiver().ok_or(TaskInitError::SetpointWatchCountTooLow)
}

static GYRO_PID_CTX: StaticCell<GyroPidContext<BoardImu>> = StaticCell::new();

/// Context for `gyro_pid` task.
#[allow(unused)]
pub struct GyroPidContext<I: ImuDevice> {
    pub imu: I,
    pub rx_receiver: RxMessageReceiver,
    pub gyro_pid_sender: GyroPidSender,
    pub setpoint_sender: SetpointSender,
    pub fast_config_subscriber: FastConfigSubscriber,
    pub imu_filters: ImuFilterBank,
    pub sensor_fusion: MadgwickFilterf32,
    pub flight_controller: FlightController,
    pub rc_controls: RcControls,
    pub rc_modes: BitSet64,
    pub gyro_pid_send_count: u32,
    pub gyro_pid_denominator: u32,
    #[cfg(feature = "autopilot")]
    pub imu_integral_sender: ImuIntegralSender,
    #[cfg(feature = "autopilot")]
    pub imu_integrator: ImuIntegrator,
}

pub fn init(
    imu: BoardImu,
    imu_filter_bank_config: ImuFilterBankConfig,
    #[cfg(feature = "rpm_filters")] rpm_notch_filter_bank_config: motor_mixers::RpmNotchFilterBankConfig,
    #[cfg(feature = "rpm_filters")] looptime_seconds: f32,
) -> Result<&'static mut GyroPidContext<BoardImu>, TaskInitError> {
    let ctx = GyroPidContext {
        imu,
        rx_receiver: rx_message_receiver()?,
        gyro_pid_sender: gyro_pid_sender(),
        #[cfg(feature = "autopilot")]
        imu_integral_sender: imu_integral_sender(),
        setpoint_sender: setpoint_sender(),
        fast_config_subscriber: fast_config_subscriber()?,

        #[cfg(feature = "rpm_filters")]
        imu_filters: ImuFilterBank::with_config_and_notch(
            imu_filter_bank_config,
            rpm_notch_filter_bank_config,
            looptime_seconds,
        ),
        #[cfg(not(feature = "rpm_filters"))]
        imu_filters: ImuFilterBank::with_config(imu_filter_bank_config),

        sensor_fusion: MadgwickFilterf32::new(),
        flight_controller: FlightController::new(),
        rc_controls: RcControls::new(),
        rc_modes: BitSet64::new(),
        gyro_pid_send_count: 0,
        gyro_pid_denominator: 10,
        #[cfg(feature = "autopilot")]
        imu_integrator: ImuIntegrator::new(),
    };

    Ok(GYRO_PID_CTX.init(ctx))
}

/// The GYRO/PID task.
/// The `gyro_pid` task calculates the motor commands, sends them immediately to the `motor_mixer` task
/// and then updates the `GyroPidMessage` and sends it, so it can be picked up the the Blackbox and the OSD.
#[embassy_executor::task]
pub async fn run(ctx: &'static mut GyroPidContext<BoardImu>) {
    log::info!("    GYRO_PID: task started");

    // Initialize IMU to use radians/s for gyro units and m/s^2 for acc units.
    let sample_rates = ctx.imu.init(8000, GyroFullScale::Max, GyroUnits::Rps, AccFullScale::Max, AccUnits::Mps2).await;
    let (gyro_rate_hz, _acc_rate_hz) = match sample_rates {
        Ok((gyro_rate_hz, acc_rate_hz)) => (gyro_rate_hz, acc_rate_hz),
        Err(_err) => (1000, 1000),
    };

    #[cfg(feature = "autopilot")]
    {
        // Set the acc_integrator denominator so that the Kalman filter runs at approximately 500 Hz
        // In practice the IMU runs at one of 1000, 2000, 4000, or 8000 Hz
        // except for the Bosch BMI270 which runs at one of 800, 1600, or 3200 Hz.
        let denominator = match gyro_rate_hz {
            32000.. => 64,
            16000..32000 => 32,
            8000..16000 => 16, // 8000 => 500Hz
            3200..8000 => 8,   // 3200 => 400Hz, 4000 => 500Hz
            1600..3200 => 4,   // 1600 => 400Hz, 2000 => 500Hz
            800..1600 => 2,    // 800 => 400Hz, 1000 => 500Hz
            ..800 => 1,
        };
        ctx.imu_integrator.set_denominator(denominator);
    }

    // TODO: this ticker wait should be replaced by an interrupt driven DMA read from the IMU.
    // I'm using a ticker like this during development to keep things simple.
    let mut ticker = embassy_time::Ticker::every(embassy_time::Duration::from_hz(u64::from(gyro_rate_hz)));

    let mut time_previous_us = embassy_time::Instant::now().as_micros();

    // This is the famous GYRO/PID loop!
    let mut loop_count: u32 = 0;
    loop {
        ticker.next().await;

        let time_now_us = embassy_time::Instant::now().as_micros();
        let delta_t_us = time_now_us - time_previous_us;
        time_previous_us = time_now_us;
        #[allow(clippy::cast_precision_loss)]
        let delta_t = (delta_t_us as f32) * 1_000_000.0;

        if let Err(err) = gyro_pid_loop_iteration(ctx, delta_t).await {
            // We really shouldn't get here, so consider so sort of failsafe handling.
            log::error!("GYRO_PID iteration failed: {err:?}");
        }
        if loop_count.is_multiple_of(1000) {
            log::info!("        GYRO_PID: loop {loop_count}");
        }
        loop_count = loop_count.wrapping_add(1);
    }
}

async fn gyro_pid_loop_iteration(ctx: &mut GyroPidContext<BoardImu>, delta_t: f32) -> Result<(), ImuError> {
    // ****
    // The GYRO part of the GYRO/PID loop
    // ****

    let (acc_unfiltered_mps2, gyro_unfiltered_rps) = ctx.imu.read_acc_gyro().await?;

    // Filter the acc and gyro values. This includes RPM notch filtering, if that is enabled.
    let (acc_mps2, gyro_rps) = ctx.imu_filters.update(acc_unfiltered_mps2, gyro_unfiltered_rps, delta_t);

    // Check if there has been a yaw heading correction from the GPS, if so, apply it.
    #[cfg(feature = "gps")]
    if let Some(message) = GPS_YAW_HEADING_SIGNAL.try_take() {
        _ = ctx.sensor_fusion.correct_yaw(message.yaw_heading_radians, message.delta_t);
    }

    // Calculate the orientation quaternion using sensor fusion.
    let orientation = ctx.sensor_fusion.fuse_acc_gyro(acc_mps2, gyro_rps, delta_t);

    #[cfg(feature = "autopilot")]
    {
        ctx.imu_integrator.integrate(acc_unfiltered_mps2, orientation, delta_t);
        // Send the integrated_imu message on a denominator
        if ctx.imu_integrator.integration_complete() {
            // Send the integrated_imu message, this will be picked up by the autopilot task.
            // It's used by the Kalman filter to estimate position and velocity.
            let imu_integral = ctx.imu_integrator.imu_integral();
            ctx.imu_integral_sender.send(ImuIntegralMessage { imu_integral });
            // Reset acc integrator for next time round the loop.
            ctx.imu_integrator.reset();

            // If the acc_bias has been updated, then set the new value before the next integration batch.
            if let Some(message) = super::autopilot::KALMAN_ACC_BIAS_SIGNAL.try_take() {
                ctx.imu_integrator.set_acc_bias(message.acc_bias_ned_mps2);
            }
        }
    }

    // ****
    // The PID part of the GYRO/PID loop
    // ****

    // If there are new control values from the radio then use them, otherwise keep using the old control values.
    if let Some(rx_message) = ctx.rx_receiver.try_changed() {
        ctx.rc_controls = rx_message.rc_controls;
        ctx.rc_modes = rx_message.rc_modes;
    }

    // Calculate the motor commands:
    // the flight controller updates its setpoints from the radio control_message
    // and then updates the PIDs using `gyro_rps` and `orientation`.
    // `setpoints_updated` is set if the setpoints have been updated because of a new radio_control_message,
    // or if the flight controller has updated the setpoints because of crash or spin recovery.
    let (motor_commands, setpoints_updated) =
        ctx.flight_controller.calculate_motor_commands(gyro_rps, orientation, delta_t, ctx.rc_controls, ctx.rc_modes);

    // Convert the motor commands calculated by the flight controller into a motor mixer message and send that message.
    // The signal will be picked up by the motor mixer task.
    // We signal every time round the GYRO/PID loop since the motor mixer also updates the RPM notch filters on this signal.
    MOTOR_MIXER_SIGNAL.signal(MotorMixerMessage::from(motor_commands));

    #[cfg(any(feature = "blackbox", feature = "osd"))]
    {
        // Send the GyroPidMessage on a denominator (eg 1/8)
        // This will be picked up by the Blackbox and the OSD.
        ctx.gyro_pid_send_count += 1;
        if ctx.gyro_pid_send_count >= ctx.gyro_pid_denominator {
            ctx.gyro_pid_send_count = 0;

            let roll_errors = ctx.flight_controller.roll.rate_pid.error();
            let pitch_errors = ctx.flight_controller.pitch.rate_pid.error();
            let yaw_errors = ctx.flight_controller.yaw_rate_pid.error();
            let pid_errors_p = [roll_errors.p, pitch_errors.p, yaw_errors.p];
            let pid_errors_i = [roll_errors.i, pitch_errors.i, yaw_errors.i];
            let pid_errors_d = [roll_errors.d, pitch_errors.d];

            let time_us = embassy_time::Instant::now().as_micros();
            let gyro_pid_message = GyroPidMessage {
                orientation,
                motor_commands,
                acc: acc_mps2,
                gyro_rps,
                gyro_rps_unfiltered: gyro_unfiltered_rps,
                pid_errors_p,
                pid_errors_i,
                pid_errors_d,
                time_us,
            };
            ctx.gyro_pid_sender.send(gyro_pid_message);

            if setpoints_updated {
                // Only send a setpoint_message when the setpoints have actually been updated.
                // This is picked up by the Blackbox.
                // TODO fill out missing SetpointMessage fields.
                let pid_errors_s = [roll_errors.s, pitch_errors.s, yaw_errors.s];
                let pid_errors_k = [roll_errors.k, pitch_errors.k, yaw_errors.k];
                let setpoint_message = SetpointMessage {
                    time_us,
                    rc_modes: ctx.rc_modes,
                    setpoints: [
                        ctx.rc_controls.roll_stick_dps,
                        ctx.rc_controls.pitch_stick_dps,
                        ctx.rc_controls.yaw_stick_dps,
                        ctx.rc_controls.throttle_stick,
                    ],
                    pid_errors_s,
                    pid_errors_k,
                    rc_commands: ctx.rc_controls.controls_pwm,

                    #[cfg(feature = "dshot_telemetry")]
                    motor_rpm_d2: [0i16; SetpointMessage::MAX_SUPPORTED_MOTOR_COUNT],
                    #[cfg(feature = "servos")]
                    servos: [0i16; SetpointMessage::MAX_SUPPORTED_SERVO_COUNT],

                    gps_state_flags: 0,
                    failsafe_phase: 0,
                    rx_signal_received: true,
                    rx_flight_channel_is_valid: true,
                };
                ctx.setpoint_sender.send(setpoint_message);
            }
        }
    }
    #[cfg(not(any(feature = "blackbox", feature = "osd")))]
    {
        // Use up unused variables.
        _ = gyro_unfiltered_rps;
        _ = setpoints_updated;
    }

    // ****
    // Check if there has been in-flight adjustment of the PID gains, if so apply them.
    // This happens infrequently.
    // ****
    //
    // try_next_message() is a simple pointer check. If there's no message, it returns None instantly,
    // so it won't mess up the 8kHz timing.
    if let Some(wait_result) = ctx.fast_config_subscriber.try_next_message()
        && let WaitResult::Message(fast_config_item) = wait_result
    {
        ctx.flight_controller.adjust_pid_gains(fast_config_item);
    }

    Ok(())
}