drone_controller/src/stacked.rs

#![allow(dead_code)]

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

use crate::input::ModeInput;
use crate::stacked::modules::ModuleRuntime;
use crate::{input::Input, DroneController};

use crate::config::*;

pub mod mixer;
pub mod modules;

use mixer::MotorMixer;
use modules::*;

#[derive(Default)]
pub struct DroneState {
    pub rotation: na::UnitQuaternion<f32>,
    pub angvel: na::Vector3<f32>,
    pub linvel: na::Vector3<f32>,
    pub position: na::Vector3<f32>,
    pub mass: f32,
}

pub struct StackedController {
    pub rotation_rt: ModuleRuntime<RotationController>,
    pub rate_rt: ModuleRuntime<AngularRateController>,
    pub linaccel_rt: ModuleRuntime<AccelerationController>,
    pub linvel_rt: ModuleRuntime<VelocityController>,

    mixer: MotorMixer,
    config: ControllerConfig,

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

impl StackedController {
    pub fn new(config: ControllerConfig) -> Self {
        let min_throttle = 0.0;
        let max_throttle = 1.0;

        let lin_vel_ctrl = VelocityController::new(PidProcessor::new(&config.stack.linvel_pid));
        let lin_accel_ctrl =
            AccelerationController::new(max_throttle, min_throttle, config.max_thrust);
        let rotation_ctrl = RotationController::new(PidProcessor::new(&config.stack.rotation_pid));
        let rate_ctrl = AngularRateController::new(PidProcessor::new(&config.stack.rate_pid));

        Self {
            linaccel_rt: ModuleRuntime::new(lin_accel_ctrl, config.stack.rotation_pid.frequency),
            rotation_rt: ModuleRuntime::new(rotation_ctrl, config.stack.rotation_pid.frequency),
            rate_rt: ModuleRuntime::new(rate_ctrl, config.stack.rate_pid.frequency),
            linvel_rt: ModuleRuntime::new(lin_vel_ctrl, config.stack.linvel_pid.frequency),
            mixer: MotorMixer {
                motor_map: config.motor_map,
                min_throttle: 0.0,
                max_throttle: 1.0,
                mixing_mode: mixer::MotorMixingMode::default(),
            },

            drone_state: DroneState {
                rotation: na::UnitQuaternion::identity(),
                mass: config.mass,
                ..Default::default()
            },
            input: Input::default(),
            last_time: 0.0,
            current_time: 0.0,

            config: config, // Keep config copy for future uses
        }
    }

    pub fn set_input(&mut self, inp: Input) {
        self.input = inp;
    }
}

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.angvel = vel;
    }

    fn set_linvel(&mut self, vel: na::Vector3<f32>) {
        self.drone_state.linvel = vel;
    }

    fn set_pos(&mut self, pos: na::Vector3<f32>) {
        self.drone_state.position = pos;
    }

    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 frame_dt = (self.current_time - self.last_time).max(0.0);
        let mut throttle = self.input.joystick.throttle_input;

        let angular_rate_setpoint = if self.input.mode != ModeInput::Acro {
            let rotation_setpoint = if self.input.mode != ModeInput::Rotation {
                let lin_accel_setpoint = if self.input.mode != ModeInput::Acceleration {
                    let lin_vel_setpoint: Velocity = if self.input.mode != ModeInput::Velocity {
                        println!("LAYER NOT IMPLENETED. TAKING 0 AS INPUT;");
                        Velocity(na::Vector3::zeros())
                    } else {
                        Velocity(na::vector![
                            self.input.velocity.x,
                            self.input.velocity.y,
                            self.input.velocity.z
                        ])
                    };
                    self.linvel_rt
                        .update(lin_vel_setpoint, &self.drone_state, frame_dt)
                } else {
                    Acceleration(na::vector![
                        self.input.acceleration.x,
                        self.input.acceleration.y,
                        self.input.acceleration.z
                    ])
                };
                let ret = self
                    .linaccel_rt
                    .update(lin_accel_setpoint, &self.drone_state, frame_dt);
                throttle = ret.1 .0;
                ret.0
            } else {
                // println!("Rotation!");
                Rotation(na::vector![
                    self.input.rotation.roll,
                    self.input.rotation.pitch,
                    self.input.rotation.yaw,
                ])
            };
            self.rotation_rt
                .update(rotation_setpoint, &self.drone_state, frame_dt)
        } else {
            AngularRate(
                na::vector![
                    self.input.joystick.roll_input,
                    self.input.joystick.pitch_input,
                    self.input.joystick.yaw_input,
                ] * self.config.stack.max_rate,
            )
        };

        let torque = self
            .rate_rt
            .update(angular_rate_setpoint, &self.drone_state, frame_dt);

        self.mixer.mix(throttle, torque.0)
    }

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