RustPhysicsMQ/src/drone/pidcontroller.rs

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

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

const PROPORTIONAL_ONLY: bool = false;

#[derive(Default)]
pub struct PIDControllerState {
    last_time: f32,
    current_time: f32,
    last_rotation: na::Unit<na::Quaternion<f32>>,
    current_rotation: na::Unit<na::Quaternion<f32>>,
    angular_velocity: na::Vector3<f32>,
    last_error: na::Vector3<f32>,
    error_sum: na::Vector3<f32>,

    pub log_time: Vec<f32>,
    pub log_target: Vec<na::Vector3<f32>>,
    pub log_current: Vec<na::Vector3<f32>>,
    pub log_error: Vec<na::Vector3<f32>>,
}

impl PIDControllerState {
    pub fn dt(&self) -> f32 {
        return self.current_time - self.last_time;
    }
}

pub struct PIDController {
    input: JoystickInput,
    pub target_rate: f32,
    pub proportional_multiplier: na::Vector3<f32>,
    pub integral_multiplier: na::Vector3<f32>,
    pub diferential_multiplier: na::Vector3<f32>,
    pub multiply_mode: bool,
    pub state: PIDControllerState,
}

impl Default for PIDController {
    fn default() -> Self {
        Self {
            input: JoystickInput::default(),
            target_rate: f32::consts::PI,
            state: Default::default(),
            multiply_mode: false,
            proportional_multiplier: vector![0.13, 3.0, 0.13],
            integral_multiplier: vector![0.0, 0.0, 0.0],
            diferential_multiplier: vector![0.001, 0.01, 0.001],
        }
    }
}

impl PIDController {
    /// Save log buffers to a CSV file
    pub fn save_logs_to_csv(&self, path: &str) -> std::io::Result<()> {
        use std::fs::File;
        use std::io::Write;

        let mut file = File::create(path)?;
        writeln!(
            file,
            "time,target_x,target_y,target_z,current_x,current_y,current_z,error_x,error_y,error_z"
        )?;

        for i in 0..self.state.log_time.len() {
            let t = self.state.log_time[i];
            let tg = self.state.log_target[i];
            let cur = self.state.log_current[i];
            let err = self.state.log_error[i];
            writeln!(
                file,
                "{},{},{},{},{},{},{},{},{},{}",
                t, tg.x, tg.y, tg.z, cur.x, cur.y, cur.z, err.x, err.y, err.z
            )?;
        }

        Ok(())
    }

    pub fn set_input(&mut self, input: JoystickInput) {
        self.input = input;
    }
    pub fn get_angular_velocity(&self) -> na::Vector3<f32> {
        return self.state.angular_velocity;
    }

    pub fn get_desired_angular_velocity(&self) -> na::Vector3<f32> {
        let coords = na::Vector3::new(
            self.input.roll_input,
            self.input.yaw_input,
            self.input.pitch_input,
        ) * self.target_rate;
        return coords;
    }
}

impl DroneController for PIDController {
    fn set_rotation(&mut self, rotation: nalgebra::Unit<nalgebra::Quaternion<f32>>) {
        self.state.last_rotation = self.state.current_rotation;
        self.state.current_rotation = rotation;
    }
    fn set_angular_velocity(&mut self, angvel: nalgebra::Vector3<f32>) {
        self.state.angular_velocity = angvel;
    }
    fn set_time(&mut self, time: f32) {
        self.state.last_time = self.state.current_time;
        self.state.current_time = time;
    }
    fn set_motor_characteristics(&self, _motor_characteristics: &MotorCharacteristics) {}
    fn get_motor_throttles(&mut self) -> [f32; 4] {
        let rot = self.state.current_rotation;
        let current = rot.inverse().transform_vector(&self.get_angular_velocity());
        let target = self.get_desired_angular_velocity();

        let error: na::Vector3<f32> = target - current;
        let throttle = self.input.throttle_input;

        self.state.log_time.push(self.state.current_time);
        self.state.log_target.push(target);
        self.state.log_current.push(current);
        self.state.log_error.push(error);

        let error_dif: na::Vector3<f32> = error - self.state.last_error;

        let forces_to_apply = match PROPORTIONAL_ONLY {
            true => {
                vector![
                    error.x * self.proportional_multiplier.x,
                    error.y * self.proportional_multiplier.y,
                    error.z * self.proportional_multiplier.z
                ]
            }
            false => {
                vector![
                    error.x * self.proportional_multiplier.x
                        + self.state.error_sum.x * self.integral_multiplier.x
                        + error_dif.x * self.diferential_multiplier.x,
                    error.y * self.proportional_multiplier.y
                        + self.state.error_sum.y * self.integral_multiplier.y
                        + error_dif.y * self.diferential_multiplier.y,
                    error.z * self.proportional_multiplier.z
                        + self.state.error_sum.z * self.integral_multiplier.z
                        + error_dif.z * self.diferential_multiplier.z,
                ]
            }
        };
        self.state.error_sum += error;
        self.state.last_error = error;

        let roll = forces_to_apply.z;
        let pitch = forces_to_apply.x;
        let yaw = forces_to_apply.y;

        let mut motors: [f32; 4] = match self.multiply_mode {
            true => {
                let mut t = [
                    1.0 + yaw - pitch + roll,
                    1.0 - yaw - pitch - roll,
                    1.0 + yaw + pitch - roll,
                    1.0 - yaw + pitch + roll,
                ];
                for m in t.iter_mut() {
                    *m *= self.input.throttle_input;
                }
                t
            }
            false => {
                let t = [
                    throttle + yaw - pitch + roll,
                    throttle - yaw - pitch - roll,
                    throttle + yaw + pitch - roll,
                    throttle - yaw + pitch + roll,
                ];
                t
            }
        };

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

        return motors;
    }

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