RustPhysicsMQ/src/drone/stacked.rs

#![allow(dead_code)]

use nalgebra::{self as na, vector};
use std::{any::Any, f32};

use crate::drone::controller::DroneController;
use crate::drone::JoystickInput;

use crate::config::*;

pub mod mixer;
pub mod modules;

use mixer::MotorMixer;
use modules::{ControllerModule, PidProcessor};

pub struct DroneState {
    pub rotation: na::UnitQuaternion<f32>,
    pub angular_velocity: na::Vector3<f32>,
}

pub struct StackedController {
    modules: Vec<ControllerModule>,
    config: SimulationConfig,
    mixer: MotorMixer,

    // State
    drone_state: DroneState,
    input: JoystickInput,
    last_time: f32,
    current_time: f32,
}

impl StackedController {
    pub fn set_input(&mut self, inp: JoystickInput) {
        self.input = inp;
    }
    pub fn new(config: SimulationConfig) -> Self {
        let mut modules = Vec::new();

        for layer in &config.layers {
            match layer {
                LayerConfig::Angle { pid, max_angle } => {
                    modules.push(ControllerModule::Angle {
                        processor: PidProcessor::new(pid),
                        max_angle: *max_angle,
                    });
                }
                LayerConfig::Rate { pid, max_rate } => {
                    modules.push(ControllerModule::Rate {
                        processor: PidProcessor::new(pid),
                        max_rate: *max_rate,
                    });
                }
            }
        }

        Self {
            mixer: MotorMixer {
                motor_map: config.motor_map,
                min_throttle: 0.0,
                max_throttle: 1.0,
            },
            modules,
            config,
            input: JoystickInput::default(),
            drone_state: DroneState {
                rotation: na::UnitQuaternion::identity(),
                angular_velocity: na::Vector3::zeros(),
            },
            last_time: 0.0,
            current_time: 0.0,
        }
    }
}

impl DroneController for StackedController {
    fn set_rotation(&mut self, rot: na::UnitQuaternion<f32>) {
        self.drone_state.rotation = rot;
    }
    fn set_angular_velocity(&mut self, vel: na::Vector3<f32>) {
        self.drone_state.angular_velocity = vel;
    }
    fn set_time(&mut self, t: f32) {
        self.last_time = self.current_time;
        self.current_time = t;
    }

    fn get_motor_throttles(&mut self) -> [f32; 4] {
        let dt = (self.current_time - self.last_time).max(0.0001);

        let mut setpoint = vector![
            self.input.roll_input,
            self.input.yaw_input,
            self.input.pitch_input,
        ];

        for (i, module) in self.modules.iter_mut().enumerate() {
            let is_first_layer = i == 0;
            setpoint = module.process(setpoint, &self.drone_state, dt, is_first_layer);
        }

        return self.mixer.mix(self.input.throttle_input, setpoint).0;
    }

    fn as_any(&self) -> &dyn Any {
        self
    }
    fn as_mut_any(&mut self) -> &mut dyn Any {
        self
    }
}