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 struct DroneState {
    pub rotation: na::UnitQuaternion<f32>,
    pub angular_velocity: na::Vector3<f32>,
}

pub enum ControllerModule {
    Rate {
        processor: PidProcessor,
        max_rate: f32,
    },
    Angle {
        processor: PidProcessor,
        max_angle: f32,
    },
}

impl ControllerModule {
    /// Takes the current setpoint (from joystick or previous layer) and
    /// returns the command for the next layer.
    pub fn process(
        &mut self,
        setpoint: na::Vector3<f32>,
        state: &DroneState,
        dt: f32,
        is_first_layer: bool,
    ) -> na::Vector3<f32> {
        match self {
            ControllerModule::Angle {
                processor,
                max_angle,
            } => {
                // Setpoint is -1.0..1.0, scale it to target Radians
                let target_angles = setpoint * *max_angle;

                let (r, p, y) = state.rotation.euler_angles();
                let current_angles = na::vector![r, y, p];

                // Output of Angle PID = Desired Angular Velocity
                processor.update(target_angles, current_angles, dt)
            }

            ControllerModule::Rate {
                processor,
                max_rate,
            } => {
                // If Rate is the start of the chain (Acro mode), scale the joystick.
                // If it's the second layer, the setpoint is already a velocity from the Angle layer.
                let target_velocity = if is_first_layer {
                    setpoint * *max_rate
                } else {
                    setpoint
                };
                let rot = state.rotation;
                let current = rot.inverse().transform_vector(&state.angular_velocity);

                // Output of Rate PID = Desired Torque/Correction Force
                processor.update(target_velocity, current, dt)
            }
        }
    }
}

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

    // 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 {
            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);

        // Input normalized -1.0 to 1.0 from joystick
        let mut setpoint = vector![
            self.input.roll_input,
            self.input.yaw_input,
            self.input.pitch_input,
        ];

        // --- CASCADED PROCESSING ---
        // If Layer 0 is Angle, setpoint becomes Desired Rate.
        // If Layer 1 is Rate, it takes that Desired Rate and outputs Torque.
        //

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

        // --- MOTOR MIXER ---
        let mut motors = [0.0; 4];
        for i in 0..4 {
            let map = self.config.motor_map[i];
            motors[i] = self.input.throttle_input
                + (map[0] * setpoint.x) // Roll
                + (map[1] * setpoint.y) // Yaw
                + (map[2] * setpoint.z); // Pitch
        }

        // Clamp while keeping relative values
        let max = motors
            .iter()
            .copied()
            .max_by(|a, b| a.total_cmp(b))
            .unwrap_or(0.0);
        if max > 1.0 {
            for v in motors.iter_mut() {
                *v /= max;
            }
        }

        motors
    }

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

pub struct PidProcessor {
    kp: na::Vector3<f32>,
    ki: na::Vector3<f32>,
    kd: na::Vector3<f32>,
    integral: na::Vector3<f32>,
    last_error: na::Vector3<f32>,
}

impl PidProcessor {
    fn new(c: &PidConfig) -> Self {
        Self {
            kp: c.p_vec(),
            ki: c.i_vec(),
            kd: c.d_vec(),
            integral: na::Vector3::zeros(),
            last_error: na::Vector3::zeros(),
        }
    }

    fn update(
        &mut self,
        target: na::Vector3<f32>,
        current: na::Vector3<f32>,
        dt: f32,
    ) -> na::Vector3<f32> {
        let error = target - current;
        self.integral += error * dt;
        let derivative = (error - self.last_error) / dt;
        self.last_error = error;

        error.component_mul(&self.kp)
            + self.integral.component_mul(&self.ki)
            + derivative.component_mul(&self.kd)
    }
}