peng_quad/
lib.rs

1//! # Peng - A Minimal Quadrotor Autonomy Framework
2//!
3//! A high-performance quadrotor autonomy framework written in Rust that provides
4//! real-time dynamics simulation, trajectory planning, and control with modern
5//! visualization capabilities.
6//!
7//! # Features
8//!
9//! ## Real-time Simulation
10//! - High-fidelity quadrotor dynamics with configurable parameters
11//! - IMU and depth sensor simulation
12//! - Optional RK4 integration for accurate dynamics
13//!
14//! ## Advanced Control
15//! - PID control for position and attitude with tunable gains
16//! - Integral windup prevention
17//! - Support for different control frequencies
18//!
19//! ## Rich Trajectory Planning
20//! - Minimum jerk line trajectory planner
21//! - Lissajous curve planner
22//! - Circular trajectory planner
23//! - Obstacle avoidance planner
24//! - Waypoint navigation planner
25//! - Landing planner
26//!
27//! ## Visualization & Debug
28//! - Real-time 3D visualization via rerun.io
29//! - Depth map rendering
30//! - State telemetry logging
31//! - Configurable logging frequencies
32//!
33//! ## Performance
34//! - Memory-safe and Efficient Rust implementation
35//! - Multi-threaded depth rendering
36//!
37//! ## Example
38//! ```
39//! use nalgebra::Vector3;
40//! use peng_quad::{Quadrotor, SimulationError};
41//! let (time_step, mass, gravity, drag_coefficient) = (0.01, 1.3, 9.81, 0.01);
42//! let inertia_matrix = [0.0347563, 0.0, 0.0, 0.0, 0.0458929, 0.0, 0.0, 0.0, 0.0977];
43//! let quadrotor = Quadrotor::new(time_step, mass, gravity, drag_coefficient, inertia_matrix);
44//! ```
45use rand::SeedableRng;
46use rayon::prelude::*;
47pub mod config;
48use nalgebra::{
49    DMatrix, DVector, Matrix3, Quaternion, Rotation3, SMatrix, UnitQuaternion, Vector3,
50};
51use osqp::{CscMatrix, Problem, Settings};
52use rand_chacha::ChaCha8Rng;
53use rand_distr::{Distribution, Normal};
54use std::f32::consts::PI;
55use std::ops::AddAssign;
56#[derive(thiserror::Error, Debug)]
57/// Represents errors that can occur during simulation
58/// # Example
59/// ```
60/// use peng_quad::SimulationError;
61/// let error = SimulationError::NalgebraError("Matrix inversion failed".to_string());
62/// ```
63pub enum SimulationError {
64    /// Error related to Rerun visualization
65    #[error("Rerun error: {0}")]
66    RerunError(#[from] rerun::RecordingStreamError),
67    /// Error related to Rerun spawn process
68    #[error("Rerun spawn error: {0}")]
69    RerunSpawnError(#[from] rerun::SpawnError),
70    /// Error related to Rerun logger setup
71    #[error("Rerun SetLogger error: {0}")]
72    SetLoggerError(#[from] rerun::external::log::SetLoggerError),
73    /// Error related to linear algebra operations
74    #[error("Nalgebra error: {0}")]
75    NalgebraError(String),
76    /// Error related to the OSQP solver
77    #[error("OSQP error: {0}")]
78    OSQPError(String),
79    /// Error related to normal distribution calculations
80    #[error("Normal error: {0}")]
81    NormalError(#[from] rand_distr::NormalError),
82    /// Other general errors
83    #[error("Other error: {0}")]
84    OtherError(String),
85}
86/// Represents a quadrotor with its physical properties and state
87/// # Example
88/// ```
89/// use nalgebra::Vector3;
90/// use peng_quad::Quadrotor;
91/// let (time_step, mass, gravity, drag_coefficient) = (0.01, 1.3, 9.81, 0.01);
92/// let inertia_matrix = [0.0347563, 0.0, 0.0, 0.0, 0.0458929, 0.0, 0.0, 0.0, 0.0977];
93/// let quadrotor = Quadrotor::new(time_step, mass, gravity, drag_coefficient, inertia_matrix);
94/// ```
95pub struct Quadrotor {
96    /// Current position of the quadrotor in 3D space
97    pub position: Vector3<f32>,
98    /// Current velocity of the quadrotor
99    pub velocity: Vector3<f32>,
100    /// Current acceleration of the quadrotor
101    pub acceleration: Vector3<f32>,
102    /// Current orientation of the quadrotor
103    pub orientation: UnitQuaternion<f32>,
104    /// Current angular velocity of the quadrotor
105    pub angular_velocity: Vector3<f32>,
106    /// Mass of the quadrotor in kg
107    pub mass: f32,
108    /// Gravitational acceleration in m/s^2
109    pub gravity: f32,
110    /// Simulation time step in seconds
111    pub time_step: f32,
112    /// Drag coefficient
113    pub drag_coefficient: f32,
114    /// Inertia matrix of the quadrotor
115    pub inertia_matrix: Matrix3<f32>,
116    /// Inverse of the inertia matrix
117    pub inertia_matrix_inv: Matrix3<f32>,
118}
119/// Implementation of the Quadrotor struct
120impl Quadrotor {
121    /// Creates a new Quadrotor with default parameters
122    /// # Arguments
123    /// * `time_step` - The simulation time step in seconds
124    /// * `mass` - The mass of the quadrotor in kg
125    /// * `gravity` - The gravitational acceleration in m/s^2
126    /// * `drag_coefficient` - The drag coefficient
127    /// * `inertia_matrix` - The inertia matrix of the quadrotor
128    /// # Returns
129    /// * A new Quadrotor instance
130    /// # Errors
131    /// * Returns a SimulationError if the inertia matrix cannot be inverted
132    /// # Example
133    /// ```
134    /// use nalgebra::Vector3;
135    /// use peng_quad::Quadrotor;
136    ///
137    /// let (time_step, mass, gravity, drag_coefficient) = (0.01, 1.3, 9.81, 0.01);
138    /// let inertia_matrix = [0.0347563, 0.0, 0.0, 0.0, 0.0458929, 0.0, 0.0, 0.0, 0.0977];
139    /// let quadrotor = Quadrotor::new(time_step, mass, gravity, drag_coefficient, inertia_matrix);
140    /// ```
141    pub fn new(
142        time_step: f32,
143        mass: f32,
144        gravity: f32,
145        drag_coefficient: f32,
146        inertia_matrix: [f32; 9],
147    ) -> Result<Self, SimulationError> {
148        let inertia_matrix = Matrix3::from_row_slice(&inertia_matrix);
149        let inertia_matrix_inv =
150            inertia_matrix
151                .try_inverse()
152                .ok_or(SimulationError::NalgebraError(
153                    "Failed to invert inertia matrix".to_string(),
154                ))?;
155        Ok(Self {
156            position: Vector3::zeros(),
157            velocity: Vector3::zeros(),
158            acceleration: Vector3::zeros(),
159            orientation: UnitQuaternion::identity(),
160            angular_velocity: Vector3::zeros(),
161            mass,
162            gravity,
163            time_step,
164            drag_coefficient,
165            inertia_matrix,
166            inertia_matrix_inv,
167        })
168    }
169    /// Updates the quadrotor's dynamics with control inputs
170    /// # Arguments
171    /// * `control_thrust` - The total thrust force applied to the quadrotor
172    /// * `control_torque` - The 3D torque vector applied to the quadrotor
173    /// # Example
174    /// ```
175    /// use nalgebra::Vector3;
176    /// use peng_quad::Quadrotor;
177    ///
178    /// let (time_step, mass, gravity, drag_coefficient) = (0.01, 1.3, 9.81, 0.01);
179    /// let inertia_matrix = [0.0347563, 0.0, 0.0, 0.0, 0.0458929, 0.0, 0.0, 0.0, 0.0977];
180    /// let mut quadrotor = Quadrotor::new(time_step, mass, gravity, drag_coefficient, inertia_matrix).unwrap();
181    /// let control_thrust = mass * gravity;
182    /// let control_torque = Vector3::new(0.0, 0.0, 0.0);
183    /// quadrotor.update_dynamics_with_controls_euler(control_thrust, &control_torque);
184    /// ```
185    pub fn update_dynamics_with_controls_euler(
186        &mut self,
187        control_thrust: f32,
188        control_torque: &Vector3<f32>,
189    ) {
190        let gravity_force = Vector3::new(0.0, 0.0, -self.mass * self.gravity);
191        let drag_force = -self.drag_coefficient * self.velocity.norm() * self.velocity;
192        let thrust_world = self.orientation * Vector3::new(0.0, 0.0, control_thrust);
193        self.acceleration = (thrust_world + gravity_force + drag_force) / self.mass;
194        self.velocity += self.acceleration * self.time_step;
195        self.position += self.velocity * self.time_step;
196        let inertia_angular_velocity = self.inertia_matrix * self.angular_velocity;
197        let gyroscopic_torque = self.angular_velocity.cross(&inertia_angular_velocity);
198        let angular_acceleration = self.inertia_matrix_inv * (control_torque - gyroscopic_torque);
199        self.angular_velocity += angular_acceleration * self.time_step;
200        self.orientation *=
201            UnitQuaternion::from_scaled_axis(self.angular_velocity * self.time_step);
202    }
203    /// Updates the quadrotor's dynamics with control inputs using the Runge-Kutta 4th order method
204    /// # Arguments
205    /// * `control_thrust` - The total thrust force applied to the quadrotor
206    /// * `control_torque` - The 3D torque vector applied to the quadrotor
207    /// # Example
208    /// ```
209    /// use nalgebra::Vector3;
210    /// use peng_quad::Quadrotor;
211    /// let (time_step, mass, gravity, drag_coefficient) = (0.01, 1.3, 9.81, 0.01);
212    /// let inertia_matrix = [0.0347563, 0.0, 0.0, 0.0, 0.0458929, 0.0, 0.0, 0.0, 0.0977];
213    /// let mut quadrotor = Quadrotor::new(time_step, mass, gravity, drag_coefficient, inertia_matrix).unwrap();
214    /// let control_thrust = mass * gravity;
215    /// let control_torque = Vector3::new(0.0, 0.0, 0.0);
216    /// quadrotor.update_dynamics_with_controls_rk4(control_thrust, &control_torque);
217    /// ```
218    pub fn update_dynamics_with_controls_rk4(
219        &mut self,
220        control_thrust: f32,
221        control_torque: &Vector3<f32>,
222    ) {
223        let h = self.time_step;
224        let state = self.get_state();
225
226        let k1 = self.state_derivative(&state, control_thrust, control_torque);
227        let mut temp_state = [0.0; 13];
228        for i in 0..13 {
229            temp_state[i] = state[i] + 0.5 * h * k1[i];
230        }
231        let k2 = self.state_derivative(&temp_state, control_thrust, control_torque);
232
233        for i in 0..13 {
234            temp_state[i] = state[i] + 0.5 * h * k2[i];
235        }
236        let k3 = self.state_derivative(&temp_state, control_thrust, control_torque);
237
238        for i in 0..13 {
239            temp_state[i] = state[i] + h * k3[i];
240        }
241        let k4 = self.state_derivative(&temp_state, control_thrust, control_torque);
242
243        let mut new_state = [0.0; 13];
244        for i in 0..13 {
245            new_state[i] = state[i] + (h / 6.0) * (k1[i] + 2.0 * k2[i] + 2.0 * k3[i] + k4[i]);
246        }
247
248        let mut q = Quaternion::new(new_state[9], new_state[6], new_state[7], new_state[8]);
249        q = q.normalize();
250        new_state[6..10].copy_from_slice(q.coords.as_slice());
251
252        self.set_state(&new_state);
253    }
254    /// Returns the state derivative of the quadrotor
255    /// # Arguments
256    /// * `state` - The state of the quadrotor
257    /// # Example
258    /// ```
259    /// use nalgebra::Vector3;
260    /// use peng_quad::Quadrotor;
261    /// use nalgebra::UnitQuaternion;
262    /// let (time_step, mass, gravity, drag_coefficient) = (0.01, 1.3, 9.81, 0.01);
263    /// let inertia_matrix = [0.0347563, 0.0, 0.0, 0.0, 0.0458929, 0.0, 0.0, 0.0, 0.0977];
264    /// let quadrotor = Quadrotor::new(time_step, mass, gravity, drag_coefficient, inertia_matrix).unwrap();
265    /// let state = quadrotor.get_state();
266    /// ```
267    pub fn get_state(&self) -> [f32; 13] {
268        let mut state = [0.0; 13];
269        state[0..3].copy_from_slice(self.position.as_slice());
270        state[3..6].copy_from_slice(self.velocity.as_slice());
271        state[6..10].copy_from_slice(self.orientation.coords.as_slice());
272        state[10..13].copy_from_slice(self.angular_velocity.as_slice());
273        state
274    }
275    /// Sets the state of the quadrotor
276    /// # Arguments
277    /// * `state` - The state of the quadrotor
278    /// # Example
279    /// ```
280    /// use nalgebra::Vector3;
281    /// use peng_quad::Quadrotor;
282    /// use nalgebra::UnitQuaternion;
283    /// let (time_step, mass, gravity, drag_coefficient) = (0.01, 1.3, 9.81, 0.01);
284    /// let inertia_matrix = [0.0347563, 0.0, 0.0, 0.0, 0.0458929, 0.0, 0.0, 0.0, 0.0977];
285    /// let mut quadrotor = Quadrotor::new(time_step, mass, gravity, drag_coefficient, inertia_matrix).unwrap();
286    /// let state = [
287    ///    0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0
288    /// ];
289    /// quadrotor.set_state(&state);
290    /// ```
291    pub fn set_state(&mut self, state: &[f32; 13]) {
292        self.position = Vector3::from_column_slice(&state[0..3]);
293        self.velocity = Vector3::from_column_slice(&state[3..6]);
294        self.orientation = UnitQuaternion::from_quaternion(Quaternion::new(
295            state[9], state[6], state[7], state[8],
296        ));
297        self.angular_velocity = Vector3::from_column_slice(&state[10..13]);
298    }
299    /// Calculates the derivative of the state of the quadrotor
300    /// # Arguments
301    /// * `state` - The current state of the quadrotor
302    /// * `control_thrust` - The thrust applied to the quadrotor
303    /// * `control_torque` - The torque applied to the quadrotor
304    /// # Returns
305    /// The derivative of the state
306    /// # Example
307    /// ```
308    /// use nalgebra::Vector3;
309    /// use peng_quad::Quadrotor;
310    /// use nalgebra::UnitQuaternion;
311    /// let (time_step, mass, gravity, drag_coefficient) = (0.01, 1.3, 9.81, 0.01);
312    /// let inertia_matrix = [0.0347563, 0.0, 0.0, 0.0, 0.0458929, 0.0, 0.0, 0.0, 0.0977];
313    /// let mut quadrotor = Quadrotor::new(time_step, mass, gravity, drag_coefficient, inertia_matrix).unwrap();
314    /// let state = [
315    ///   0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0
316    /// ];
317    /// let control_thrust = 0.0;
318    /// let control_torque = Vector3::new(0.0, 0.0, 0.0);
319    /// let derivative = quadrotor.state_derivative(&state, control_thrust, &control_torque);
320    /// ```
321    pub fn state_derivative(
322        &mut self,
323        state: &[f32],
324        control_thrust: f32,
325        control_torque: &Vector3<f32>,
326    ) -> [f32; 13] {
327        let velocity = Vector3::from_column_slice(&state[3..6]);
328        let orientation = UnitQuaternion::from_quaternion(Quaternion::new(
329            state[9], state[6], state[7], state[8],
330        ));
331        // Quaternion deriviative
332        let omega_quat = Quaternion::new(0.0, state[10], state[11], state[12]);
333        let q_dot = orientation.into_inner() * omega_quat * 0.5;
334
335        let angular_velocity = Vector3::from_column_slice(&state[10..13]);
336
337        let gravity_force = Vector3::new(0.0, 0.0, -self.mass * self.gravity);
338        let drag_force = -self.drag_coefficient * velocity.norm() * velocity;
339        let thrust_world = orientation * Vector3::new(0.0, 0.0, control_thrust);
340        self.acceleration = (thrust_world + gravity_force + drag_force) / self.mass;
341
342        let inertia_angular_velocity = self.inertia_matrix * angular_velocity;
343        let gyroscopic_torque = angular_velocity.cross(&inertia_angular_velocity);
344        let angular_acceleration = self.inertia_matrix_inv * (control_torque - gyroscopic_torque);
345
346        let mut derivative = [0.0; 13];
347        derivative[0..3].copy_from_slice(velocity.as_slice());
348        derivative[3..6].copy_from_slice(self.acceleration.as_slice());
349        derivative[6..10].copy_from_slice(q_dot.coords.as_slice());
350        derivative[10..13].copy_from_slice(angular_acceleration.as_slice());
351        derivative
352    }
353    /// Simulates IMU readings
354    /// # Returns
355    /// * A tuple containing the true acceleration and angular velocity of the quadrotor
356    /// # Errors
357    /// * Returns a SimulationError if the IMU readings cannot be calculated
358    /// # Example
359    /// ```
360    /// use nalgebra::Vector3;
361    /// use peng_quad::Quadrotor;
362    ///
363    /// let (time_step, mass, gravity, drag_coefficient) = (0.01, 1.3, 9.81, 0.01);
364    /// let inertia_matrix = [0.0347563, 0.0, 0.0, 0.0, 0.0458929, 0.0, 0.0, 0.0, 0.0977];
365    /// let quadrotor = Quadrotor::new(time_step, mass, gravity, drag_coefficient, inertia_matrix).unwrap();
366    /// let (true_acceleration, true_angular_velocity) = quadrotor.read_imu().unwrap();
367    /// ```
368    pub fn read_imu(&self) -> Result<(Vector3<f32>, Vector3<f32>), SimulationError> {
369        Ok((self.acceleration, self.angular_velocity))
370    }
371}
372/// Represents an Inertial Measurement Unit (IMU) with bias and noise characteristics
373/// # Example
374/// ```
375/// use nalgebra::Vector3;
376/// use peng_quad::Imu;
377/// let accel_noise_std = 0.0003;
378/// let gyro_noise_std = 0.02;
379/// let accel_bias_std = 0.0001;
380/// let gyro_bias_std = 0.001;
381/// let imu = Imu::new(accel_noise_std, gyro_noise_std, accel_bias_std, gyro_bias_std);
382/// ```
383pub struct Imu {
384    /// Accelerometer bias
385    pub accel_bias: Vector3<f32>,
386    /// Gyroscope bias
387    pub gyro_bias: Vector3<f32>,
388    /// Standard deviation of accelerometer noise
389    pub accel_noise_std: f32,
390    /// Standard deviation of gyroscope noise
391    pub gyro_noise_std: f32,
392    /// Standard deviation of accelerometer bias drift
393    pub accel_bias_std: f32,
394    /// Standard deviation of gyroscope bias drift
395    pub gyro_bias_std: f32,
396    /// Accelerometer noise distribution
397    accel_noise: Normal<f32>,
398    /// Gyroscope noise distribution
399    gyro_noise: Normal<f32>,
400    /// Accelerometer bias drift distribution
401    accel_bias_drift: Normal<f32>,
402    /// Gyroscope bias drift distribution
403    gyro_bias_drift: Normal<f32>,
404    /// Random number generator
405    rng: ChaCha8Rng,
406}
407/// Implements the IMU
408impl Imu {
409    /// Creates a new IMU with default parameters
410    /// # Arguments
411    /// * `accel_noise_std` - Standard deviation of accelerometer noise
412    /// * `gyro_noise_std` - Standard deviation of gyroscope noise
413    /// * `accel_bias_std` - Standard deviation of accelerometer bias drift
414    /// * `gyro_bias_std` - Standard deviation of gyroscope bias drift
415    /// # Returns
416    /// * A new Imu instance
417    /// # Example
418    /// ```
419    /// use peng_quad::Imu;
420    ///
421    /// let imu = Imu::new(0.01, 0.01, 0.01, 0.01);
422    /// ```
423    pub fn new(
424        accel_noise_std: f32,
425        gyro_noise_std: f32,
426        accel_bias_std: f32,
427        gyro_bias_std: f32,
428    ) -> Result<Self, SimulationError> {
429        Ok(Self {
430            accel_bias: Vector3::zeros(),
431            gyro_bias: Vector3::zeros(),
432            accel_noise_std,
433            gyro_noise_std,
434            accel_bias_std,
435            gyro_bias_std,
436            accel_noise: Normal::new(0.0, accel_noise_std)?,
437            gyro_noise: Normal::new(0.0, gyro_noise_std)?,
438            accel_bias_drift: Normal::new(0.0, accel_bias_std)?,
439            gyro_bias_drift: Normal::new(0.0, gyro_bias_std)?,
440            rng: ChaCha8Rng::from_os_rng(),
441        })
442    }
443    /// Updates the IMU biases over time
444    /// # Arguments
445    /// * `dt` - Time step for the update
446    /// # Errors
447    /// * Returns a SimulationError if the bias drift cannot be calculated
448    /// # Example
449    /// ```
450    /// use peng_quad::Imu;
451    ///
452    /// let mut imu = Imu::new(0.01, 0.01, 0.01, 0.01).unwrap();
453    /// imu.update(0.01).unwrap();
454    /// ```
455    pub fn update(&mut self, dt: f32) -> Result<(), SimulationError> {
456        let dt_sqrt = fast_sqrt(dt);
457        let accel_drift = self.accel_bias_drift.sample(&mut self.rng) * dt_sqrt;
458        let gyro_drift = self.gyro_bias_drift.sample(&mut self.rng) * dt_sqrt;
459        self.accel_bias += Vector3::from_iterator((0..3).map(|_| accel_drift));
460        self.gyro_bias += Vector3::from_iterator((0..3).map(|_| gyro_drift));
461        Ok(())
462    }
463    /// Simulates IMU readings with added bias and noise
464    ///
465    /// The added bias and noise are based on normal distributions
466    /// # Arguments
467    /// * `true_acceleration` - The true acceleration vector
468    /// * `true_angular_velocity` - The true angular velocity vector
469    /// # Returns
470    /// * A tuple containing the measured acceleration and angular velocity
471    /// # Errors
472    /// * Returns a SimulationError if the IMU readings cannot be calculated
473    /// # Example
474    /// ```
475    /// use nalgebra::Vector3;
476    /// use peng_quad::Imu;
477    ///
478    /// let mut imu = Imu::new(0.01, 0.01, 0.01, 0.01).unwrap();
479    /// let true_acceleration = Vector3::new(0.0, 0.0, 9.81);
480    /// let true_angular_velocity = Vector3::new(0.0, 0.0, 0.0);
481    /// let (measured_acceleration, measured_ang_velocity) = imu.read(true_acceleration, true_angular_velocity).unwrap();
482    /// ```
483    pub fn read(
484        &mut self,
485        true_acceleration: Vector3<f32>,
486        true_angular_velocity: Vector3<f32>,
487    ) -> Result<(Vector3<f32>, Vector3<f32>), SimulationError> {
488        let accel_noise_sample =
489            Vector3::from_iterator((0..3).map(|_| self.accel_noise.sample(&mut self.rng)));
490        let gyro_noise_sample =
491            Vector3::from_iterator((0..3).map(|_| self.gyro_noise.sample(&mut self.rng)));
492        let measured_acceleration = true_acceleration + self.accel_bias + accel_noise_sample;
493        let measured_ang_velocity = true_angular_velocity + self.gyro_bias + gyro_noise_sample;
494        Ok((measured_acceleration, measured_ang_velocity))
495    }
496}
497/// PID controller for quadrotor position and attitude control
498///
499/// The kpid_pos and kpid_att gains are following the format of
500/// proportional, derivative and integral gains
501/// # Example
502/// ```
503/// use nalgebra::Vector3;
504/// use peng_quad::PIDController;
505/// let kpid_pos = [
506///     [1.0, 1.0, 1.0],
507///     [0.1, 0.1, 0.1],
508///     [0.01, 0.01, 0.01],
509/// ];
510/// let kpid_att = [
511///     [1.0, 1.0, 1.0],
512///     [0.1, 0.1, 0.1],
513///     [0.01, 0.01, 0.01],
514/// ];
515/// let max_integral_pos = [1.0, 1.0, 1.0];
516/// let max_integral_att = [1.0, 1.0, 1.0];
517/// let mass = 1.0;
518/// let gravity = 9.81;
519/// let pid_controller = PIDController::new(kpid_pos, kpid_att, max_integral_pos, max_integral_att, mass, gravity);
520/// ```
521pub struct PIDController {
522    /// PID gain for position control including proportional, derivative, and integral gains
523    pub kpid_pos: [Vector3<f32>; 3],
524    /// PID gain for attitude control including proportional, derivative, and integral gains
525    pub kpid_att: [Vector3<f32>; 3],
526    /// Accumulated integral error for position
527    pub integral_pos_error: Vector3<f32>,
528    /// Accumulated integral error for attitude
529    pub integral_att_error: Vector3<f32>,
530    /// Maximum allowed integral error for position
531    pub max_integral_pos: Vector3<f32>,
532    /// Maximum allowed integral error for attitude
533    pub max_integral_att: Vector3<f32>,
534    /// Mass of the quadrotor
535    pub mass: f32,
536    /// Gravity constant
537    pub gravity: f32,
538}
539/// Implementation of PIDController
540impl PIDController {
541    /// Creates a new PIDController with default gains
542    /// gains are in the order of proportional, derivative, and integral
543    /// # Arguments
544    /// * `_kpid_pos` - PID gains for position control
545    /// * `_kpid_att` - PID gains for attitude control
546    /// * `_max_integral_pos` - Maximum allowed integral error for position
547    /// * `_max_integral_att` - Maximum allowed integral error for attitude
548    /// # Returns
549    /// * A new PIDController instance
550    /// # Example
551    /// ```
552    /// use nalgebra::Vector3;
553    /// use peng_quad::PIDController;
554    /// let kpid_pos = [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]];
555    /// let kpid_att = [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]];
556    /// let max_integral_pos = [1.0, 1.0, 1.0];
557    /// let max_integral_att = [1.0, 1.0, 1.0];
558    /// let mass = 1.0;
559    /// let gravity = 9.81;
560    /// let pid = PIDController::new(kpid_pos, kpid_att, max_integral_pos, max_integral_att, mass, gravity);
561    /// ```
562    pub fn new(
563        _kpid_pos: [[f32; 3]; 3],
564        _kpid_att: [[f32; 3]; 3],
565        _max_integral_pos: [f32; 3],
566        _max_integral_att: [f32; 3],
567        _mass: f32,
568        _gravity: f32,
569    ) -> Self {
570        Self {
571            kpid_pos: _kpid_pos.map(Vector3::from),
572            kpid_att: _kpid_att.map(Vector3::from),
573            integral_pos_error: Vector3::zeros(),
574            integral_att_error: Vector3::zeros(),
575            max_integral_pos: Vector3::from(_max_integral_pos),
576            max_integral_att: Vector3::from(_max_integral_att),
577            mass: _mass,
578            gravity: _gravity,
579        }
580    }
581    /// Computes attitude control torques
582    /// # Arguments
583    /// * `desired_orientation` - The desired orientation quaternion
584    /// * `current_orientation` - The current orientation quaternion
585    /// * `current_angular_velocity` - The current angular velocity
586    /// * `dt` - Time step
587    /// # Returns
588    /// * The computed attitude control torques
589    /// # Example
590    /// ```
591    /// use nalgebra::{UnitQuaternion, Vector3};
592    /// use peng_quad::PIDController;
593    ///
594    /// let kpid_pos = [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]];
595    /// let kpid_att = [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]];
596    /// let max_integral_pos = [1.0, 1.0, 1.0];
597    /// let max_integral_att = [1.0, 1.0, 1.0];
598    /// let mass = 1.0;
599    /// let gravity = 9.81;
600    /// let mut pid = PIDController::new(kpid_pos, kpid_att, max_integral_pos, max_integral_att, mass, gravity);
601    /// let desired_orientation = UnitQuaternion::identity();
602    /// let current_orientation = UnitQuaternion::identity();
603    /// let current_angular_velocity = Vector3::zeros();
604    /// let dt = 0.01;
605    /// let control_torques = pid.compute_attitude_control(&desired_orientation, &current_orientation, &current_angular_velocity, dt);
606    /// ```
607    pub fn compute_attitude_control(
608        &mut self,
609        desired_orientation: &UnitQuaternion<f32>,
610        current_orientation: &UnitQuaternion<f32>,
611        current_angular_velocity: &Vector3<f32>,
612        dt: f32,
613    ) -> Vector3<f32> {
614        let error_orientation = current_orientation.inverse() * desired_orientation;
615        let (roll_error, pitch_error, yaw_error) = error_orientation.euler_angles();
616        let error_angles = Vector3::new(roll_error, pitch_error, yaw_error);
617        self.integral_att_error += error_angles * dt;
618        self.integral_att_error = self
619            .integral_att_error
620            .zip_map(&self.max_integral_att, |int, max| int.clamp(-max, max));
621        let error_angular_velocity = -current_angular_velocity; // TODO: Add desired angular velocity
622        self.kpid_att[0].component_mul(&error_angles)
623            + self.kpid_att[1].component_mul(&error_angular_velocity)
624            + self.kpid_att[2].component_mul(&self.integral_att_error)
625    }
626    /// Computes position control thrust and desired orientation
627    /// # Arguments
628    /// * `desired_position` - The desired position
629    /// * `desired_velocity` - The desired velocity
630    /// * `desired_yaw` - The desired yaw angle
631    /// * `current_position` - The current position
632    /// * `current_velocity` - The current velocity
633    /// * `dt` - Time step
634    /// * `mass` - Mass of the quadrotor
635    /// * `gravity` - Gravitational acceleration
636    /// # Returns
637    /// * A tuple containing the computed thrust and desired orientation quaternion
638    /// # Example
639    /// ```
640    /// use nalgebra::{UnitQuaternion, Vector3};
641    /// use peng_quad::PIDController;
642    ///
643    /// let kpid_pos = [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]];
644    /// let kpid_att = [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]];
645    /// let max_integral_pos = [1.0, 1.0, 1.0];
646    /// let max_integral_att = [1.0, 1.0, 1.0];
647    /// let mass = 1.0;
648    /// let gravity = 9.81;
649    /// let mut pid = PIDController::new(kpid_pos, kpid_att, max_integral_pos, max_integral_att, mass, gravity);
650    /// let desired_position = Vector3::new(0.0, 0.0, 1.0);
651    /// let desired_velocity = Vector3::zeros();
652    /// let desired_yaw = 0.0;
653    /// let current_position = Vector3::zeros();
654    /// let current_velocity = Vector3::zeros();
655    /// let dt = 0.01;
656    /// let (thrust, desired_orientation) = pid.compute_position_control(&desired_position, &desired_velocity, desired_yaw, &current_position, &current_velocity, dt);
657    /// ```
658    pub fn compute_position_control(
659        &mut self,
660        desired_position: &Vector3<f32>,
661        desired_velocity: &Vector3<f32>,
662        desired_yaw: f32,
663        current_position: &Vector3<f32>,
664        current_velocity: &Vector3<f32>,
665        dt: f32,
666    ) -> (f32, UnitQuaternion<f32>) {
667        let error_position = desired_position - current_position;
668        let error_velocity = desired_velocity - current_velocity;
669        self.integral_pos_error += error_position * dt;
670        self.integral_pos_error = self
671            .integral_pos_error
672            .zip_map(&self.max_integral_pos, |int, max| int.clamp(-max, max));
673        let acceleration = self.kpid_pos[0].component_mul(&error_position)
674            + self.kpid_pos[1].component_mul(&error_velocity)
675            + self.kpid_pos[2].component_mul(&self.integral_pos_error);
676        let gravity_compensation = Vector3::new(0.0, 0.0, self.gravity);
677        let total_acceleration = acceleration + gravity_compensation;
678        let total_acc_norm = total_acceleration.norm();
679        let thrust = self.mass * total_acc_norm;
680        let desired_orientation = if total_acc_norm > 1e-6 {
681            let z_body = total_acceleration / total_acc_norm;
682            let yaw_rotation = UnitQuaternion::from_euler_angles(0.0, 0.0, desired_yaw);
683            let x_body_horizontal = yaw_rotation * Vector3::new(1.0, 0.0, 0.0);
684            let y_body = z_body.cross(&x_body_horizontal).normalize();
685            let x_body = y_body.cross(&z_body);
686            UnitQuaternion::from_rotation_matrix(&Rotation3::from_matrix_unchecked(
687                Matrix3::from_columns(&[x_body, y_body, z_body]),
688            ))
689        } else {
690            UnitQuaternion::from_euler_angles(0.0, 0.0, desired_yaw)
691        };
692        (thrust, desired_orientation)
693    }
694}
695/// Enum representing different types of trajectory planners
696/// # Example
697/// ```
698/// use peng_quad::PlannerType;
699/// use peng_quad::HoverPlanner;
700/// let hover_planner : PlannerType = PlannerType::Hover(HoverPlanner{
701///     target_position: nalgebra::Vector3::new(0.0, 0.0, 1.0),
702///     target_yaw: 0.0,
703/// });
704/// ```
705pub enum PlannerType {
706    /// Hover planner
707    Hover(HoverPlanner),
708    /// Minimum jerk line planner
709    MinimumJerkLine(MinimumJerkLinePlanner),
710    /// Minimum jerk circle planner
711    Lissajous(LissajousPlanner),
712    /// Minimum jerk circle planner
713    Circle(CirclePlanner),
714    /// Minimum jerk landing planner
715    Landing(LandingPlanner),
716    /// Obstacle avoidance planner
717    ObstacleAvoidance(ObstacleAvoidancePlanner),
718    /// Minimum snap waypoint planner
719    MinimumSnapWaypoint(MinimumSnapWaypointPlanner),
720    /// Quadratic Polynomial based waypoint planner
721    QPpolyTraj(QPpolyTrajPlanner),
722}
723/// Implementation of the planner type
724impl PlannerType {
725    /// Plans the trajectory based on the current planner type
726    /// # Arguments
727    /// * `current_position` - The current position of the quadrotor
728    /// * `current_velocity` - The current velocity of the quadrotor
729    /// * `time` - The current simulation time
730    /// # Returns
731    /// * A tuple containing the desired position, velocity, and yaw angle
732    /// # Example
733    /// ```
734    /// use nalgebra::Vector3;
735    /// use peng_quad::PlannerType;
736    /// use peng_quad::HoverPlanner;
737    /// let hover_planner = HoverPlanner {
738    ///     target_position: Vector3::new(0.0, 0.0, 1.0),
739    ///     target_yaw: 0.0
740    /// };
741    /// let hover_planner_type = PlannerType::Hover(hover_planner);
742    /// let (desired_position, desired_velocity, desired_yaw) = hover_planner_type.plan(Vector3::new(0.0, 0.0, 0.0), Vector3::new(0.0, 0.0, 0.0), 0.0);
743    /// ```
744    pub fn plan(
745        &self,
746        current_position: Vector3<f32>,
747        current_velocity: Vector3<f32>,
748        time: f32,
749    ) -> (Vector3<f32>, Vector3<f32>, f32) {
750        match self {
751            PlannerType::Hover(p) => p.plan(current_position, current_velocity, time),
752            PlannerType::MinimumJerkLine(p) => p.plan(current_position, current_velocity, time),
753            PlannerType::Lissajous(p) => p.plan(current_position, current_velocity, time),
754            PlannerType::Circle(p) => p.plan(current_position, current_velocity, time),
755            PlannerType::Landing(p) => p.plan(current_position, current_velocity, time),
756            PlannerType::ObstacleAvoidance(p) => p.plan(current_position, current_velocity, time),
757            PlannerType::MinimumSnapWaypoint(p) => p.plan(current_position, current_velocity, time),
758            PlannerType::QPpolyTraj(p) => p.plan(current_position, current_velocity, time),
759        }
760    }
761    /// Checks if the current trajectory is finished
762    /// # Arguments
763    /// * `current_position` - The current position of the quadrotor
764    /// * `time` - The current simulation time
765    /// # Returns
766    /// * `true` if the trajectory is finished, `false` otherwise
767    /// # Example
768    /// ```
769    /// use nalgebra::Vector3;
770    /// use peng_quad::PlannerType;
771    /// use peng_quad::HoverPlanner;
772    /// use peng_quad::Planner;
773    /// let hover_planner = HoverPlanner{
774    ///     target_position: Vector3::new(0.0, 0.0, 1.0),
775    ///     target_yaw: 0.0,
776    /// };
777    /// let is_finished = hover_planner.is_finished(Vector3::new(0.0, 0.0, 0.0), 0.0);
778    /// ```
779    pub fn is_finished(
780        &self,
781        current_position: Vector3<f32>,
782        time: f32,
783    ) -> Result<bool, SimulationError> {
784        match self {
785            PlannerType::Hover(p) => p.is_finished(current_position, time),
786            PlannerType::MinimumJerkLine(p) => p.is_finished(current_position, time),
787            PlannerType::Lissajous(p) => p.is_finished(current_position, time),
788            PlannerType::Circle(p) => p.is_finished(current_position, time),
789            PlannerType::Landing(p) => p.is_finished(current_position, time),
790            PlannerType::ObstacleAvoidance(p) => p.is_finished(current_position, time),
791            PlannerType::MinimumSnapWaypoint(p) => p.is_finished(current_position, time),
792            PlannerType::QPpolyTraj(p) => p.is_finished(current_position, time),
793        }
794    }
795}
796/// Trait defining the interface for trajectory planners
797/// # Example
798/// ```
799/// use nalgebra::Vector3;
800/// use peng_quad::{Planner, SimulationError};
801/// struct TestPlanner;
802/// impl Planner for TestPlanner {
803///    fn plan(
804///         &self,
805///         current_position: Vector3<f32>,
806///         current_velocity: Vector3<f32>,
807///         time: f32,
808/// ) -> (Vector3<f32>, Vector3<f32>, f32) {
809///         (Vector3::new(0.0, 0.0, 0.0), Vector3::new(0.0, 0.0, 0.0), 0.0)
810///     }
811///     fn is_finished(
812///         &self,
813///         current_position: Vector3<f32>,
814///         time: f32,
815///     ) -> Result<bool, SimulationError> {
816///         Ok(true)
817///     }
818/// }
819/// ```
820pub trait Planner {
821    /// Plans the trajectory based on the current state and time
822    /// # Arguments
823    /// * `current_position` - The current position of the quadrotor
824    /// * `current_velocity` - The current velocity of the quadrotor
825    /// * `time` - The current simulation time
826    /// # Returns
827    /// * A tuple containing the desired position, velocity, and yaw angle
828    /// # Example
829    /// ```
830    /// use nalgebra::Vector3;
831    /// use peng_quad::{Planner, SimulationError};
832    /// struct TestPlanner;
833    /// impl Planner for TestPlanner {
834    ///     fn plan(
835    ///         &self,
836    ///         current_position: Vector3<f32>,
837    ///         current_velocity: Vector3<f32>,
838    ///         time: f32,
839    /// ) -> (Vector3<f32>, Vector3<f32>, f32) {
840    ///         (Vector3::new(0.0, 0.0, 0.0), Vector3::new(0.0, 0.0, 0.0), 0.0)
841    ///     }
842    ///     fn is_finished(
843    ///         &self,
844    ///         current_position: Vector3<f32>,
845    ///         time: f32,
846    ///     ) -> Result<bool, SimulationError> {
847    ///         Ok(true)
848    ///     }
849    /// }
850    /// ```
851    fn plan(
852        &self,
853        current_position: Vector3<f32>,
854        current_velocity: Vector3<f32>,
855        time: f32,
856    ) -> (Vector3<f32>, Vector3<f32>, f32);
857    /// Checks if the current trajectory is finished
858    /// # Arguments
859    /// * `current_position` - The current position of the quadrotor
860    /// * `time` - The current simulation time
861    /// # Returns
862    /// * `true` if the trajectory is finished, `false` otherwise
863    /// # Example
864    /// ```
865    /// use nalgebra::Vector3;
866    /// use peng_quad::{Planner, SimulationError};
867    /// struct TestPlanner;
868    /// impl Planner for TestPlanner {
869    ///     fn plan(
870    ///         &self,
871    ///         current_position: Vector3<f32>,
872    ///         current_velocity: Vector3<f32>,
873    ///         time: f32,
874    /// ) -> (Vector3<f32>, Vector3<f32>, f32) {
875    ///         (Vector3::new(0.0, 0.0, 0.0), Vector3::new(0.0, 0.0, 0.0), 0.0)
876    ///     }
877    ///     fn is_finished(
878    ///         &self,
879    ///         current_position: Vector3<f32>,
880    ///         time: f32,
881    ///     ) -> Result<bool, SimulationError> {
882    ///         Ok(true)
883    ///     }
884    /// }
885    /// ```
886    fn is_finished(
887        &self,
888        current_position: Vector3<f32>,
889        time: f32,
890    ) -> Result<bool, SimulationError>;
891}
892/// Planner for hovering at a fixed position
893/// # Example
894/// ```
895/// use nalgebra::Vector3;
896/// use peng_quad::HoverPlanner;
897/// let hover_planner = HoverPlanner {
898///     target_position: Vector3::new(0.0, 0.0, 0.0),
899///     target_yaw: 0.0,
900/// };
901/// ```
902pub struct HoverPlanner {
903    /// Target position for hovering
904    pub target_position: Vector3<f32>,
905    /// Target yaw angle for hovering
906    pub target_yaw: f32,
907}
908/// Implementation of the `Planner` trait for the `HoverPlanner`
909impl Planner for HoverPlanner {
910    fn plan(
911        &self,
912        _current_position: Vector3<f32>,
913        _current_velocity: Vector3<f32>,
914        _time: f32,
915    ) -> (Vector3<f32>, Vector3<f32>, f32) {
916        (self.target_position, Vector3::zeros(), self.target_yaw)
917    }
918
919    fn is_finished(
920        &self,
921        _current_position: Vector3<f32>,
922        _time: f32,
923    ) -> Result<bool, SimulationError> {
924        Ok(false) // Hover planner never "finished"
925    }
926}
927/// Planner for minimum jerk trajectories along a straight line
928/// # Example
929/// ```
930/// use nalgebra::Vector3;
931/// use peng_quad::MinimumJerkLinePlanner;
932/// let minimum_jerk_line_planner = MinimumJerkLinePlanner {
933///     start_position: Vector3::new(0.0, 0.0, 0.0),
934///     end_position: Vector3::new(1.0, 1.0, 1.0),
935///     start_yaw: 0.0,
936///     end_yaw: 0.0,
937///     start_time: 0.0,
938///     duration: 1.0,
939/// };
940/// ```
941pub struct MinimumJerkLinePlanner {
942    /// Starting position of the trajectory
943    pub start_position: Vector3<f32>,
944    /// Ending position of the trajectory
945    pub end_position: Vector3<f32>,
946    /// Starting yaw angle
947    pub start_yaw: f32,
948    /// Ending yaw angle
949    pub end_yaw: f32,
950    /// Start time of the trajectory
951    pub start_time: f32,
952    /// Duration of the trajectory
953    pub duration: f32,
954}
955/// Implementation of the planner trait for minimum jerk line planner
956impl Planner for MinimumJerkLinePlanner {
957    fn plan(
958        &self,
959        _current_position: Vector3<f32>,
960        _current_velocity: Vector3<f32>,
961        time: f32,
962    ) -> (Vector3<f32>, Vector3<f32>, f32) {
963        let t = ((time - self.start_time) / self.duration).clamp(0.0, 1.0);
964        let t2 = t * t;
965        let t3 = t2 * t;
966        let t4 = t3 * t;
967        let s = 10.0 * t2 - 15.0 * t3 + 6.0 * t4;
968        let s_dot = (30.0 * t.powi(2) - 60.0 * t.powi(3) + 30.0 * t.powi(4)) / self.duration;
969        let position = self.start_position + (self.end_position - self.start_position) * s;
970        let velocity = (self.end_position - self.start_position) * s_dot;
971        let yaw = self.start_yaw + (self.end_yaw - self.start_yaw) * s;
972        (position, velocity, yaw)
973    }
974
975    fn is_finished(
976        &self,
977        _current_position: Vector3<f32>,
978        _time: f32,
979    ) -> Result<bool, SimulationError> {
980        Ok((_current_position - self.end_position).norm() < 0.01
981            && _time >= self.start_time + self.duration)
982    }
983}
984/// Planner for Lissajous curve trajectories
985/// # Example
986/// ```
987/// use nalgebra::Vector3;
988/// use peng_quad::LissajousPlanner;
989/// let lissajous_planner = LissajousPlanner {
990///     start_position: Vector3::new(0.0, 0.0, 0.0),
991///     center: Vector3::new(1.0, 1.0, 1.0),
992///     amplitude: Vector3::new(1.0, 1.0, 1.0),
993///     frequency: Vector3::new(1.0, 1.0, 1.0),
994///     phase: Vector3::new(0.0, 0.0, 0.0),
995///     start_time: 0.0,
996///     duration: 1.0,
997///     start_yaw: 0.0,
998///     end_yaw: 0.0,
999///     ramp_time: 0.1,
1000/// };
1001/// ```
1002pub struct LissajousPlanner {
1003    /// Starting position of the trajectory
1004    pub start_position: Vector3<f32>,
1005    /// Center of the Lissajous curve
1006    pub center: Vector3<f32>,
1007    /// Amplitude of the Lissajous curve
1008    pub amplitude: Vector3<f32>,
1009    /// Frequency of the Lissajous curve
1010    pub frequency: Vector3<f32>,
1011    /// Phase of the Lissajous curve
1012    pub phase: Vector3<f32>,
1013    /// Start time of the trajectory
1014    pub start_time: f32,
1015    /// Duration of the trajectory
1016    pub duration: f32,
1017    /// Starting yaw angle
1018    pub start_yaw: f32,
1019    /// Ending yaw angle
1020    pub end_yaw: f32,
1021    /// Ramp-up time for smooth transitions
1022    pub ramp_time: f32,
1023}
1024/// Implementation of the planner trait for Lissajous curve trajectories
1025impl Planner for LissajousPlanner {
1026    fn plan(
1027        &self,
1028        _current_position: Vector3<f32>,
1029        _current_velocity: Vector3<f32>,
1030        time: f32,
1031    ) -> (Vector3<f32>, Vector3<f32>, f32) {
1032        let t = ((time - self.start_time) / self.duration).clamp(0.0, 1.0);
1033        let smooth_start = if t < self.ramp_time / self.duration {
1034            let t_ramp = t / (self.ramp_time / self.duration);
1035            t_ramp * t_ramp * (3.0 - 2.0 * t_ramp)
1036        } else {
1037            1.0
1038        };
1039        let velocity_ramp = if t < self.ramp_time / self.duration {
1040            smooth_start
1041        } else if t > 1.0 - self.ramp_time / self.duration {
1042            let t_down = (1.0 - t) / (self.ramp_time / self.duration);
1043            t_down * t_down * (3.0 - 2.0 * t_down)
1044        } else {
1045            1.0
1046        };
1047        let ang_pos = self.frequency * t * 2.0 * PI + self.phase;
1048        let lissajous = self.amplitude.component_mul(&ang_pos.map(f32::sin));
1049        let position =
1050            self.start_position + smooth_start * ((self.center + lissajous) - self.start_position);
1051        let mut velocity = Vector3::new(
1052            self.amplitude.x * self.frequency.x * 2.0 * PI * ang_pos.x.cos(),
1053            self.amplitude.y * self.frequency.y * 2.0 * PI * ang_pos.y.cos(),
1054            self.amplitude.z * self.frequency.z * 2.0 * PI * ang_pos.z.cos(),
1055        ) * velocity_ramp
1056            / self.duration;
1057        if t < self.ramp_time / self.duration {
1058            let transition_velocity = (self.center - self.start_position)
1059                * (2.0 * t / self.ramp_time - 2.0 * t * t / (self.ramp_time * self.ramp_time))
1060                / self.duration;
1061            velocity += transition_velocity;
1062        }
1063        let yaw = self.start_yaw + (self.end_yaw - self.start_yaw) * t;
1064        (position, velocity, yaw)
1065    }
1066
1067    fn is_finished(
1068        &self,
1069        _current_position: Vector3<f32>,
1070        time: f32,
1071    ) -> Result<bool, SimulationError> {
1072        Ok(time >= self.start_time + self.duration)
1073    }
1074}
1075/// Planner for circular trajectories
1076/// # Example
1077/// ```
1078/// use nalgebra::Vector3;
1079/// use peng_quad::CirclePlanner;
1080/// let circle_planner = CirclePlanner {
1081///     center: Vector3::new(1.0, 1.0, 1.0),
1082///     radius: 1.0,
1083///     angular_velocity: 1.0,
1084///     start_position: Vector3::new(0.0, 0.0, 0.0),
1085///     start_time: 0.0,
1086///     duration: 1.0,
1087///     start_yaw: 0.0,
1088///     end_yaw: 0.0,
1089///     ramp_time: 0.1,
1090/// };
1091/// ```
1092pub struct CirclePlanner {
1093    /// Center of the circular trajectory
1094    pub center: Vector3<f32>,
1095    /// Radius of the circular trajectory
1096    pub radius: f32,
1097    /// Angular velocity of the circular motion
1098    pub angular_velocity: f32,
1099    /// Starting position of the trajectory
1100    pub start_position: Vector3<f32>,
1101    /// Start time of the trajectory
1102    pub start_time: f32,
1103    /// Duration of the trajectory
1104    pub duration: f32,
1105    /// Starting yaw angle
1106    pub start_yaw: f32,
1107    /// Ending yaw angle
1108    pub end_yaw: f32,
1109    /// Ramp-up time for smooth transitions
1110    pub ramp_time: f32,
1111}
1112/// Implementation of the Planner trait for CirclePlanner
1113impl Planner for CirclePlanner {
1114    fn plan(
1115        &self,
1116        _current_position: Vector3<f32>,
1117        _current_velocity: Vector3<f32>,
1118        time: f32,
1119    ) -> (Vector3<f32>, Vector3<f32>, f32) {
1120        let t = (time - self.start_time) / self.duration;
1121        let t = t.clamp(0.0, 1.0);
1122        let smooth_start = if t < self.ramp_time / self.duration {
1123            let t_ramp = t / (self.ramp_time / self.duration);
1124            t_ramp * t_ramp * (3.0 - 2.0 * t_ramp)
1125        } else {
1126            1.0
1127        };
1128        let velocity_ramp = if t < self.ramp_time / self.duration {
1129            smooth_start
1130        } else if t > 1.0 - self.ramp_time / self.duration {
1131            let t_down = (1.0 - t) / (self.ramp_time / self.duration);
1132            t_down * t_down * (3.0 - 2.0 * t_down)
1133        } else {
1134            1.0
1135        };
1136        let angle = self.angular_velocity * t * self.duration;
1137        let circle_offset = Vector3::new(self.radius * angle.cos(), self.radius * angle.sin(), 0.0);
1138        let position = self.start_position
1139            + smooth_start * ((self.center + circle_offset) - self.start_position);
1140        let tangential_velocity = Vector3::new(
1141            -self.radius * self.angular_velocity * angle.sin(),
1142            self.radius * self.angular_velocity * angle.cos(),
1143            0.0,
1144        );
1145        let velocity = tangential_velocity * velocity_ramp;
1146        let yaw = self.start_yaw + (self.end_yaw - self.start_yaw) * t;
1147        (position, velocity, yaw)
1148    }
1149
1150    fn is_finished(
1151        &self,
1152        _current_position: Vector3<f32>,
1153        time: f32,
1154    ) -> Result<bool, SimulationError> {
1155        Ok(time >= self.start_time + self.duration)
1156    }
1157}
1158/// Planner for landing maneuvers
1159/// # Example
1160/// ```
1161/// use nalgebra::Vector3;
1162/// use peng_quad::LandingPlanner;
1163/// let landing_planner = LandingPlanner {
1164///    start_position: Vector3::new(0.0, 0.0, 1.0),
1165///     start_time: 0.0,
1166///     duration: 1.0,
1167///     start_yaw: 0.0,
1168/// };
1169/// ```
1170pub struct LandingPlanner {
1171    /// Starting position of the landing maneuver
1172    pub start_position: Vector3<f32>,
1173    /// Start time of the landing maneuver
1174    pub start_time: f32,
1175    /// Duration of the landing maneuver
1176    pub duration: f32,
1177    /// Starting yaw angle
1178    pub start_yaw: f32,
1179}
1180/// Implementation of the Planner trait for 'LandingPlanner'
1181impl Planner for LandingPlanner {
1182    fn plan(
1183        &self,
1184        _current_position: Vector3<f32>,
1185        _current_velocity: Vector3<f32>,
1186        time: f32,
1187    ) -> (Vector3<f32>, Vector3<f32>, f32) {
1188        let t = ((time - self.start_time) / self.duration).clamp(0.0, 1.0);
1189        let target_z = self.start_position.z * (1.0 - t);
1190        let target_position = Vector3::new(self.start_position.x, self.start_position.y, target_z);
1191        let target_velocity = Vector3::new(0.0, 0.0, -self.start_position.z / self.duration);
1192        (target_position, target_velocity, self.start_yaw)
1193    }
1194
1195    fn is_finished(
1196        &self,
1197        current_position: Vector3<f32>,
1198        time: f32,
1199    ) -> Result<bool, SimulationError> {
1200        Ok(current_position.z < 0.05 || time >= self.start_time + self.duration)
1201    }
1202}
1203/// Manages different trajectory planners and switches between them
1204/// # Example
1205/// ```
1206/// use nalgebra::Vector3;
1207/// use peng_quad::PlannerManager;
1208/// let initial_position = Vector3::new(0.0, 0.0, 1.0);
1209/// let initial_yaw = 0.0;
1210/// let planner_manager = PlannerManager::new(initial_position, initial_yaw);
1211/// ```
1212pub struct PlannerManager {
1213    /// The current planner
1214    pub current_planner: PlannerType,
1215}
1216/// Implementation of the PlannerManager
1217impl PlannerManager {
1218    /// Creates a new PlannerManager with an initial hover planner
1219    /// # Arguments
1220    /// * `initial_position` - The initial position for hovering
1221    /// * `initial_yaw` - The initial yaw angle for hovering
1222    /// # Returns
1223    /// * A new PlannerManager instance
1224    /// # Example
1225    /// ```
1226    /// use nalgebra::Vector3;
1227    /// use peng_quad::PlannerManager;
1228    /// let initial_position = Vector3::new(0.0, 0.0, 1.0);
1229    /// let initial_yaw = 0.0;
1230    /// let planner_manager = PlannerManager::new(initial_position, initial_yaw);
1231    /// ```
1232    pub fn new(initial_position: Vector3<f32>, initial_yaw: f32) -> Self {
1233        let hover_planner = HoverPlanner {
1234            target_position: initial_position,
1235            target_yaw: initial_yaw,
1236        };
1237        Self {
1238            current_planner: PlannerType::Hover(hover_planner),
1239        }
1240    }
1241    /// Sets a new planner
1242    /// # Arguments
1243    /// * `new_planner` - The new planner to be set
1244    /// # Example
1245    /// ```
1246    /// use nalgebra::Vector3;
1247    /// use peng_quad::{PlannerManager, CirclePlanner, PlannerType};
1248    /// let initial_position = Vector3::new(0.0, 0.0, 1.0);
1249    /// let initial_yaw = 0.0;
1250    /// let mut planner_manager = PlannerManager::new(initial_position, initial_yaw);
1251    /// let new_planner = CirclePlanner {
1252    ///     center: Vector3::new(0.0, 0.0, 1.0),
1253    ///     radius: 1.0,
1254    ///     angular_velocity: 1.0,
1255    ///     start_yaw: 0.0,
1256    ///     end_yaw: 0.0,
1257    ///     start_time: 0.0,
1258    ///     duration: 10.0,
1259    ///     ramp_time: 1.0,
1260    ///     start_position: Vector3::new(0.0, 0.0, 1.0),
1261    /// };
1262    /// planner_manager.set_planner(PlannerType::Circle(new_planner));
1263    /// ```
1264    pub fn set_planner(&mut self, new_planner: PlannerType) {
1265        self.current_planner = new_planner;
1266    }
1267    /// Updates the current planner and returns the desired position, velocity, and yaw
1268    /// # Arguments
1269    /// * `current_position` - The current position of the quadrotor
1270    /// * `current_orientation` - The current orientation of the quadrotor
1271    /// * `current_velocity` - The current velocity of the quadrotor
1272    /// * `time` - The current simulation time
1273    /// # Returns
1274    /// * A tuple containing the desired position, velocity, and yaw angle
1275    /// # Errors
1276    /// * Returns a SimulationError if the current planner is not finished
1277    /// # Example
1278    /// ```
1279    /// use nalgebra::{Vector3, UnitQuaternion};
1280    /// use peng_quad::{PlannerManager, SimulationError};
1281    /// let initial_position = Vector3::new(0.0, 0.0, 1.0);
1282    /// let initial_yaw = 0.0;
1283    /// let mut planner_manager = PlannerManager::new(initial_position, initial_yaw);
1284    /// let current_position = Vector3::new(0.0, 0.0, 1.0);
1285    /// let current_orientation = UnitQuaternion::from_euler_angles(0.0, 0.0, 0.0);
1286    /// let current_velocity = Vector3::new(0.0, 0.0, 0.0);
1287    /// let obstacles = vec![];
1288    /// let time = 0.0;
1289    /// let result = planner_manager.update(current_position, current_orientation, current_velocity, time, &obstacles);
1290    /// match result {
1291    ///     Ok((target_position, target_velocity, target_yaw)) => {
1292    ///         println!("Target Position: {:?}", target_position);
1293    ///         println!("Target Velocity: {:?}", target_velocity);
1294    ///         println!("Target Yaw: {:?}", target_yaw);
1295    ///     }
1296    ///     Err(SimulationError) => {
1297    ///         log::error!("Error: Planner is not finished");
1298    ///     }
1299    /// }
1300    /// ```
1301    pub fn update(
1302        &mut self,
1303        current_position: Vector3<f32>,
1304        current_orientation: UnitQuaternion<f32>,
1305        current_velocity: Vector3<f32>,
1306        time: f32,
1307        obstacles: &[Obstacle],
1308    ) -> Result<(Vector3<f32>, Vector3<f32>, f32), SimulationError> {
1309        if self.current_planner.is_finished(current_position, time)? {
1310            log::info!("Time: {time:.2} s,\tSwitch Hover");
1311            self.current_planner = PlannerType::Hover(HoverPlanner {
1312                target_position: current_position,
1313                target_yaw: current_orientation.euler_angles().2,
1314            });
1315        }
1316        // Update obstacles for ObstacleAvoidancePlanner if needed
1317        if let PlannerType::ObstacleAvoidance(ref mut planner) = self.current_planner {
1318            planner.obstacles = obstacles.to_owned();
1319        }
1320        Ok(self
1321            .current_planner
1322            .plan(current_position, current_velocity, time))
1323    }
1324}
1325/// Obstacle avoidance planner that uses a potential field approach to avoid obstacles
1326///
1327/// The planner calculates a repulsive force for each obstacle and an attractive force towards the goal
1328/// The resulting force is then used to calculate the desired position and velocity
1329/// # Example
1330/// ```
1331/// use nalgebra::Vector3;
1332/// use peng_quad::{ObstacleAvoidancePlanner, Obstacle};
1333/// let planner = ObstacleAvoidancePlanner {
1334///     target_position: Vector3::new(0.0, 0.0, 1.0),
1335///     start_time: 0.0,
1336///     duration: 10.0,
1337///     start_yaw: 0.0,
1338///     end_yaw: 0.0,
1339///     obstacles: vec![Obstacle {
1340///         position: Vector3::new(1.0, 0.0, 1.0),
1341///         velocity: Vector3::new(0.0, 0.0, 0.0),
1342///         radius: 0.5,
1343///     }],
1344///     k_att: 1.0,
1345///     k_rep: 1.0,
1346///     k_vortex: 1.0,
1347///     d0: 1.0,
1348///     d_target: 1.0,
1349///     max_speed: 1.0,
1350/// };
1351/// ```
1352pub struct ObstacleAvoidancePlanner {
1353    /// Target position of the planner
1354    pub target_position: Vector3<f32>,
1355    /// Start time of the planner
1356    pub start_time: f32,
1357    /// Duration of the planner
1358    pub duration: f32,
1359    /// Starting yaw angle
1360    pub start_yaw: f32,
1361    /// Ending yaw angle
1362    pub end_yaw: f32,
1363    /// List of obstacles
1364    pub obstacles: Vec<Obstacle>,
1365    /// Attractive force gain
1366    pub k_att: f32,
1367    /// Repulsive force gain
1368    pub k_rep: f32,
1369    /// Vortex force gain
1370    pub k_vortex: f32,
1371    /// Influence distance of obstacles
1372    pub d0: f32,
1373    /// Influence distance of target
1374    pub d_target: f32,
1375    /// Maximum speed of the quadrotor
1376    pub max_speed: f32,
1377}
1378/// Implementation of the Planner trait for ObstacleAvoidancePlanner
1379impl Planner for ObstacleAvoidancePlanner {
1380    fn plan(
1381        &self,
1382        current_position: Vector3<f32>,
1383        current_velocity: Vector3<f32>,
1384        time: f32,
1385    ) -> (Vector3<f32>, Vector3<f32>, f32) {
1386        let t = ((time - self.start_time) / self.duration).clamp(0.0, 1.0);
1387        let distance_to_target = (self.target_position - current_position).norm();
1388        let f_att = self.k_att
1389            * self.smooth_attractive_force(distance_to_target)
1390            * (self.target_position - current_position).normalize();
1391        // Repulsive force from obstacles
1392        let mut f_rep = Vector3::zeros();
1393        let mut f_vortex = Vector3::zeros();
1394        for obstacle in &self.obstacles {
1395            let diff = current_position - obstacle.position;
1396            let distance = diff.norm();
1397            if distance < self.d0 {
1398                f_rep += self.k_rep
1399                    * (1.0 / distance - 1.0 / self.d0)
1400                    * (1.0 / distance.powi(2))
1401                    * diff.normalize();
1402                f_vortex +=
1403                    self.k_vortex * current_velocity.cross(&diff).normalize() / distance.powi(2);
1404            }
1405        }
1406        let f_total = f_att + f_rep + f_vortex;
1407        let desired_velocity = f_total.normalize() * self.max_speed.min(f_total.norm());
1408        let desired_position = current_position + desired_velocity * self.duration * (1.0 - t);
1409        let desired_yaw = self.start_yaw + (self.end_yaw - self.start_yaw) * t;
1410        (desired_position, desired_velocity, desired_yaw)
1411    }
1412
1413    fn is_finished(
1414        &self,
1415        current_position: Vector3<f32>,
1416        time: f32,
1417    ) -> Result<bool, SimulationError> {
1418        Ok((current_position - self.target_position).norm() < 0.1
1419            && time >= self.start_time + self.duration)
1420    }
1421}
1422/// Implementation of the ObstacleAvoidancePlanner
1423impl ObstacleAvoidancePlanner {
1424    /// A smooth attractive force function that transitions from linear to exponential decay
1425    /// When the distance to the target is less than the target distance, the force is linear
1426    /// When the distance is greater, the force decays exponentially
1427    /// # Arguments
1428    /// * `distance` - The distance to the target
1429    /// # Returns
1430    /// * The attractive force
1431    /// # Example
1432    /// ```
1433    /// use peng_quad::ObstacleAvoidancePlanner;
1434    /// let planner = ObstacleAvoidancePlanner {
1435    ///    target_position: nalgebra::Vector3::new(0.0, 0.0, 1.0),
1436    ///     start_time: 0.0,
1437    ///     duration: 10.0,
1438    ///     start_yaw: 0.0,
1439    ///     end_yaw: 0.0,
1440    ///     obstacles: vec![],
1441    ///     k_att: 1.0,
1442    ///     k_rep: 1.0,
1443    ///     k_vortex: 1.0,
1444    ///     d0: 1.0,
1445    ///     d_target: 1.0,
1446    ///     max_speed: 1.0,
1447    /// };
1448    /// let distance = 1.0;
1449    /// let force = planner.smooth_attractive_force(distance);
1450    /// ```
1451    #[inline]
1452    pub fn smooth_attractive_force(&self, distance: f32) -> f32 {
1453        if distance <= self.d_target {
1454            distance
1455        } else {
1456            self.d_target + (distance - self.d_target).tanh()
1457        }
1458    }
1459}
1460/// Waypoint planner that generates a minimum snap trajectory between waypoints
1461/// # Example
1462/// ```
1463/// use peng_quad::MinimumSnapWaypointPlanner;
1464/// use nalgebra::Vector3;
1465/// let planner = MinimumSnapWaypointPlanner::new(
1466///     vec![Vector3::new(0.0, 0.0, 0.0), Vector3::new(1.0, 0.0, 0.0)],
1467///     vec![0.0, 0.0],
1468///     vec![1.0],
1469///     0.0,
1470/// );
1471/// ```
1472pub struct MinimumSnapWaypointPlanner {
1473    /// List of waypoints
1474    pub waypoints: Vec<Vector3<f32>>,
1475    /// List of yaw angles
1476    pub yaws: Vec<f32>,
1477    /// List of segment times to reach each waypoint
1478    pub times: Vec<f32>,
1479    /// Coefficients for the x, y, and z components of the trajectory
1480    pub coefficients: Vec<Vec<Vector3<f32>>>,
1481    /// Coefficients for the yaw component of the trajectory
1482    pub yaw_coefficients: Vec<Vec<f32>>,
1483    /// Start time of the trajectory
1484    pub start_time: f32,
1485}
1486/// Implementation of the MinimumSnapWaypointPlanner
1487impl MinimumSnapWaypointPlanner {
1488    /// Generate a new minimum snap waypoint planner
1489    /// # Arguments
1490    /// * `waypoints` - List of waypoints
1491    /// * `yaws` - List of yaw angles
1492    /// * `segment_times` - List of segment times to reach each waypoint
1493    /// * `start_time` - Start time of the trajectory
1494    /// # Returns
1495    /// * A new minimum snap waypoint planner
1496    /// # Errors
1497    /// * Returns an error if the number of waypoints, yaws, and segment times do not match
1498    /// # Example
1499    /// ```
1500    /// use peng_quad::MinimumSnapWaypointPlanner;
1501    /// use nalgebra::Vector3;
1502    /// let waypoints = vec![Vector3::zeros(), Vector3::new(1.0, 0.0, 0.0)];
1503    /// let yaws = vec![0.0, 0.0];
1504    /// let segment_times = vec![1.0];
1505    /// let start_time = 0.0;
1506    /// let planner = MinimumSnapWaypointPlanner::new(waypoints, yaws, segment_times, start_time);
1507    /// ```
1508    pub fn new(
1509        waypoints: Vec<Vector3<f32>>,
1510        yaws: Vec<f32>,
1511        segment_times: Vec<f32>,
1512        start_time: f32,
1513    ) -> Result<Self, SimulationError> {
1514        if waypoints.len() < 2 {
1515            return Err(SimulationError::OtherError(
1516                "At least two waypoints are required".to_string(),
1517            ));
1518        }
1519        if waypoints.len() != segment_times.len() + 1 || waypoints.len() != yaws.len() {
1520            return Err(SimulationError::OtherError("Number of segment times must be one less than number of waypoints, and yaws must match waypoints".to_string()));
1521        }
1522        let mut planner = Self {
1523            waypoints,
1524            yaws,
1525            times: segment_times,
1526            coefficients: Vec::new(),
1527            yaw_coefficients: Vec::new(),
1528            start_time,
1529        };
1530        planner.compute_minimum_snap_trajectories()?;
1531        planner.compute_minimum_acceleration_yaw_trajectories()?;
1532        Ok(planner)
1533    }
1534    /// Compute the coefficients for the minimum snap trajectory, calculated for each segment between waypoints
1535    /// # Errors
1536    /// * Returns an error if the nalgebra solver fails to solve the linear system
1537    /// # Example
1538    /// ```
1539    /// use peng_quad::MinimumSnapWaypointPlanner;
1540    /// use nalgebra::Vector3;
1541    /// let waypoints = vec![Vector3::zeros(), Vector3::new(1.0, 0.0, 0.0)];
1542    /// let yaws = vec![0.0, 0.0];
1543    /// let segment_times = vec![1.0];
1544    /// let start_time = 0.0;
1545    /// let mut planner = MinimumSnapWaypointPlanner::new(waypoints, yaws, segment_times, start_time).unwrap();
1546    /// planner.compute_minimum_snap_trajectories();
1547    /// ```
1548    pub fn compute_minimum_snap_trajectories(&mut self) -> Result<(), SimulationError> {
1549        let n = self.waypoints.len() - 1;
1550        for i in 0..n {
1551            let duration = self.times[i];
1552            let (start, end) = (self.waypoints[i], self.waypoints[i + 1]);
1553            let mut a = SMatrix::<f32, 8, 8>::zeros();
1554            let mut b = SMatrix::<f32, 8, 3>::zeros();
1555            a.fixed_view_mut::<4, 4>(0, 0).fill_with_identity();
1556            b.fixed_view_mut::<1, 3>(0, 0).copy_from(&start.transpose());
1557            b.fixed_view_mut::<1, 3>(4, 0).copy_from(&end.transpose());
1558            // End point constraints
1559            for j in 0..8 {
1560                a[(4, j)] = duration.powi(j as i32);
1561                if j > 0 {
1562                    a[(5, j)] = j as f32 * duration.powi(j as i32 - 1);
1563                }
1564                if j > 1 {
1565                    a[(6, j)] = j as f32 * (j - 1) as f32 * duration.powi(j as i32 - 2);
1566                }
1567                if j > 2 {
1568                    a[(7, j)] =
1569                        j as f32 * (j - 1) as f32 * (j - 2) as f32 * duration.powi(j as i32 - 3);
1570                }
1571            }
1572            let coeffs = a.lu().solve(&b).ok_or(SimulationError::NalgebraError(
1573                "Failed to solve for coefficients in MinimumSnapWaypointPlanner".to_string(),
1574            ))?;
1575            self.coefficients.push(
1576                (0..8)
1577                    .map(|j| Vector3::new(coeffs[(j, 0)], coeffs[(j, 1)], coeffs[(j, 2)]))
1578                    .collect(),
1579            );
1580        }
1581        Ok(())
1582    }
1583    /// Compute the coefficients for yaw trajectories
1584    /// The yaw trajectory is a cubic polynomial and interpolated between waypoints
1585    /// # Errors
1586    /// * Returns an error if nalgebra fails to solve for the coefficients
1587    /// # Example
1588    /// ```
1589    /// use peng_quad::MinimumSnapWaypointPlanner;
1590    /// use nalgebra::Vector3;
1591    /// let waypoints = vec![Vector3::zeros(), Vector3::new(1.0, 0.0, 0.0)];
1592    /// let yaws = vec![0.0, 0.0];
1593    /// let segment_times = vec![1.0];
1594    /// let start_time = 0.0;
1595    /// let mut planner = MinimumSnapWaypointPlanner::new(waypoints, yaws, segment_times, start_time).unwrap();
1596    /// planner.compute_minimum_snap_trajectories();
1597    /// planner.compute_minimum_acceleration_yaw_trajectories();
1598    /// ```
1599    pub fn compute_minimum_acceleration_yaw_trajectories(&mut self) -> Result<(), SimulationError> {
1600        let n = self.yaws.len() - 1; // Number of segments
1601        for i in 0..n {
1602            let (duration, start_yaw, end_yaw) = (self.times[i], self.yaws[i], self.yaws[i + 1]);
1603            let mut a = SMatrix::<f32, 4, 4>::zeros();
1604            let mut b = SMatrix::<f32, 4, 1>::zeros();
1605            (a[(0, 0)], a[(1, 1)]) = (1.0, 1.0);
1606            (b[0], b[2]) = (start_yaw, end_yaw);
1607            for j in 0..4 {
1608                a[(2, j)] = duration.powi(j as i32);
1609                if j > 0 {
1610                    a[(3, j)] = j as f32 * duration.powi(j as i32 - 1);
1611                }
1612            }
1613            let yaw_coeffs = a.lu().solve(&b).ok_or(SimulationError::NalgebraError(
1614                "Failed to solve for yaw coefficients in MinimumSnapWaypointPlanner".to_string(),
1615            ))?;
1616            self.yaw_coefficients.push(yaw_coeffs.as_slice().to_vec());
1617        }
1618        Ok(())
1619    }
1620    /// Evaluate the trajectory at a given time, returns the position, velocity, yaw, and yaw rate at the given time
1621    /// # Arguments
1622    /// * `t` - The time to evaluate the trajectory at
1623    /// * `coeffs` - The coefficients for the position trajectory
1624    /// * `yaw_coeffs` - The coefficients for the yaw trajectory
1625    /// # Returns
1626    /// * `position` - The position at the given time (meters)
1627    /// * `velocity` - The velocity at the given time (meters / second)
1628    /// * `yaw` - The yaw at the given time (radians)
1629    /// * `yaw_rate` - The yaw rate at the given time (radians / second)
1630    /// # Example
1631    /// ```
1632    /// use nalgebra::Vector3;
1633    /// use peng_quad::MinimumSnapWaypointPlanner;
1634    /// let waypoints = vec![Vector3::zeros(), Vector3::new(1.0, 0.0, 0.0)];
1635    /// let yaws = vec![0.0, 0.0];
1636    /// let segment_times = vec![1.0];
1637    /// let start_time = 0.0;
1638    /// let mut planner = MinimumSnapWaypointPlanner::new(waypoints, yaws, segment_times, start_time).unwrap();
1639    /// planner.compute_minimum_snap_trajectories();
1640    /// planner.compute_minimum_acceleration_yaw_trajectories();
1641    /// let (position, velocity, yaw, yaw_rate) = planner.evaluate_polynomial(0.5, &planner.coefficients[0], &planner.yaw_coefficients[0]);
1642    /// ```
1643    pub fn evaluate_polynomial(
1644        &self,
1645        t: f32,
1646        coeffs: &[Vector3<f32>],
1647        yaw_coeffs: &[f32],
1648    ) -> (Vector3<f32>, Vector3<f32>, f32, f32) {
1649        let mut position = Vector3::zeros();
1650        let mut velocity = Vector3::zeros();
1651        let mut yaw = 0.0;
1652        let mut yaw_rate = 0.0;
1653        for (i, coeff) in coeffs.iter().enumerate() {
1654            let ti = t.powi(i as i32);
1655            position += coeff * ti;
1656            if i > 0 {
1657                velocity += coeff * (i as f32) * t.powi(i as i32 - 1);
1658            }
1659        }
1660        for (i, &coeff) in yaw_coeffs.iter().enumerate() {
1661            let ti = t.powi(i as i32);
1662            yaw += coeff * ti;
1663            if i > 0 {
1664                yaw_rate += coeff * (i as f32) * t.powi(i as i32 - 1);
1665            }
1666        }
1667        (position, velocity, yaw, yaw_rate)
1668    }
1669}
1670/// Implement the `Planner` trait for `MinimumSnapWaypointPlanner`
1671impl Planner for MinimumSnapWaypointPlanner {
1672    fn plan(
1673        &self,
1674        _current_position: Vector3<f32>,
1675        _current_velocity: Vector3<f32>,
1676        time: f32,
1677    ) -> (Vector3<f32>, Vector3<f32>, f32) {
1678        let relative_time = time - self.start_time;
1679        // Find the current segment
1680        let mut segment_start_time = 0.0;
1681        let mut current_segment = 0;
1682        for (i, &segment_duration) in self.times.iter().enumerate() {
1683            if relative_time < segment_start_time + segment_duration {
1684                current_segment = i;
1685                break;
1686            }
1687            segment_start_time += segment_duration;
1688        }
1689        // Evaluate the polynomial for the current segment
1690        let segment_time = relative_time - segment_start_time;
1691        let (position, velocity, yaw, _yaw_rate) = self.evaluate_polynomial(
1692            segment_time,
1693            &self.coefficients[current_segment],
1694            &self.yaw_coefficients[current_segment],
1695        );
1696        (position, velocity, yaw)
1697    }
1698
1699    fn is_finished(
1700        &self,
1701        current_position: Vector3<f32>,
1702        time: f32,
1703    ) -> Result<bool, SimulationError> {
1704        let last_waypoint = self.waypoints.last().ok_or(SimulationError::OtherError(
1705            "No waypoints available".to_string(),
1706        ))?;
1707        Ok(time >= self.start_time + self.times.iter().sum::<f32>()
1708            && (current_position - last_waypoint).norm() < 0.1)
1709    }
1710}
1711
1712/// Waypoint planner that generates a quadratic polynomial trajectory between waypoints
1713/// # Example
1714/// ```
1715/// use peng_quad::QPpolyTrajPlanner;
1716/// use nalgebra::Vector3;
1717/// let waypoints: Vec<Vec<f32>> = vec![vec![0.0,0.0,1.0,0.0], vec![1.0,0.0,1.0,0.0]];
1718/// let segment_times = vec![6.0];
1719/// let min_deriv = 3;
1720/// let polyorder = 9;
1721/// let smooth_upto = 4;
1722/// let max_velocity = 4.0;
1723/// let max_acceleration = 3.0;
1724/// let dt = 0.1;
1725/// let start_time = 0.0;
1726/// let mut qp_planner = QPpolyTrajPlanner::new(waypoints,segment_times,polyorder, min_deriv, smooth_upto,max_velocity,max_acceleration, start_time, dt);
1727/// ```
1728pub struct QPpolyTrajPlanner {
1729    /// Matrix of coefficients for each segment and each dimension, organized as nrows: polyorder*segment_times.len(), ncols: 4 (for x, y, z, yaw)
1730    pub coeff: DMatrix<f64>,
1731    /// Order of the polynomial to be used in computing trajectory
1732    pub polyorder: usize,
1733    /// Minimize which derivative in the QP problem (1->Velocity, 2->Acceleration, 3->Snap, 4->Jerk. Please note that derivative value greater than 4 is not supported)
1734    pub min_deriv: usize,
1735    /// Ensure continuity upto which derivative. NOTE: This MUST be <= polynomial_order (1->Velocity, 2->Acceleration, 3->Snap, 4->Jerk. Please note that derivative value greater than 4 is not supported)
1736    pub smooth_upto: usize,
1737    /// Vector of time values for each segment, which tells the planner how much time each segment should take to complete. Expressed in seconds.
1738    pub segment_times: Vec<f32>,
1739    /// Waypoints for each segment. Note that there should be segment_times.len() + 1 values for waypoints, with the position of the quadrotor being the very first waypoint.
1740    pub waypoints: Vec<Vec<f32>>,
1741    /// Maximum velocity constraint. Set to 0.0 to disregard inequality constraints. Please set reasonable values for this as it influences solver convergence or failure.
1742    pub max_velocity: f32,
1743    /// Maximum acceleration constraint. Set to 0.0 to disregard inequality constraints. Please set reasonable values for this as it influences solver convergence or failure.
1744    pub max_acceleration: f32,
1745    /// Time at which the simulation starts. This value has no bearing to QPpolyTraj itself, but is used during simulation since we use relative time internally when computing values.
1746    pub start_time: f32,
1747    /// Step time used while generating inequality constraints. Has no bearing if max_velocity or max_acceleration is set to 0.0. Please set reasonable value for this as it has a huge impact on OSQP solve time.
1748    pub dt: f32,
1749}
1750
1751impl QPpolyTrajPlanner {
1752    /// Generate a new QPpolyTraj planner
1753    /// # Arguments
1754    /// * `waypoints` - The waypoints for the trajectory
1755    /// * `segment_times` - The times for each segment of the trajectory
1756    /// * `polyorder` - The order of the polynomial
1757    /// * `min_deriv` - The minimum derivative to be considered
1758    /// * `smooth_upto` - The maximum derivative to be considered
1759    /// * `max_velocity` - The maximum velocity
1760    /// * `max_acceleration` - The maximum acceleration
1761    /// * `start_time` - The start time of the trajectory
1762    /// * `dt` - The time step for the trajectory
1763    /// # Errors
1764    /// * Returns an error if the specified waypoints, segment times, smooth_upto or polyorder is invalid
1765    /// # Examples
1766    /// ```
1767    /// use peng_quad::QPpolyTrajPlanner;
1768    /// use nalgebra::{DMatrix, DVector, Vector3};
1769    /// let waypoints: Vec<Vec<f32>> = vec![vec![0.0, 0.0, 0.0, 0.0], vec![1.0, 1.0, 1.5, 1.5707963267948966], vec![-1.0, 1.0, 1.75, 3.141592653589793], vec![0.0, -1.0, 1.0, -1.5707963267948966], vec![0.0, 0.0, 0.5, 0.0]];
1770    /// let segment_times = vec![5.0, 5.0, 5.0, 5.0];
1771    /// let polyorder = 9;
1772    /// let min_deriv = 3;
1773    /// let smooth_upto = 4;
1774    /// let max_velocity = 4.0;
1775    /// let max_acceleration = 2.0;
1776    /// let start_time = 0.0;
1777    /// let dt = 0.1;
1778    /// let mut qp_planner = QPpolyTrajPlanner::new(waypoints,segment_times,polyorder, min_deriv, smooth_upto,max_velocity,max_acceleration, start_time, dt).unwrap();
1779    /// ```
1780    #[allow(clippy::too_many_arguments)]
1781    pub fn new(
1782        waypoints: Vec<Vec<f32>>,
1783        segment_times: Vec<f32>,
1784        polyorder: usize,
1785        min_deriv: usize,
1786        smooth_upto: usize,
1787        max_velocity: f32,
1788        max_acceleration: f32,
1789        start_time: f32,
1790        dt: f32,
1791    ) -> Result<Self, SimulationError> {
1792        if waypoints.len() < 2 {
1793            return Err(SimulationError::OtherError(
1794                "At least two waypoints are required".to_string(),
1795            ));
1796        }
1797        if waypoints.len() != segment_times.len() + 1 || waypoints[0].len() != 4 {
1798            return Err(SimulationError::OtherError("Number of segment times must be one less than number of waypoints, and waypoints must contain x, y, z, yaw values".to_string()));
1799        }
1800        if smooth_upto > polyorder {
1801            return Err(SimulationError::OtherError(format!(
1802                "smooth_upto ({smooth_upto}) cannot be greater than polynomial order ({polyorder})"
1803            )));
1804        }
1805        let mut planner = Self {
1806            coeff: DMatrix::zeros(segment_times.len() * polyorder, 4),
1807            polyorder,
1808            min_deriv,
1809            smooth_upto,
1810            segment_times,
1811            waypoints,
1812            max_velocity,
1813            max_acceleration,
1814            start_time,
1815            dt,
1816        };
1817        planner.compute_coefficients()?;
1818        Ok(planner)
1819    }
1820
1821    /// Compute the position, velocity, yaw, and yaw_rate for a given time t in given segment
1822    /// # Arguments
1823    /// * `t` - The time at which to evaluate the polynomial
1824    /// * `segment` - The segment index for which to evaluate the polynomial
1825    /// # Returns
1826    /// * A tuple containing the position, velocity, yaw, and yaw_rate at time `t` in segment `segment`
1827    pub fn evaluate_polynomial(
1828        &self,
1829        t: f32,
1830        segment: usize,
1831    ) -> (Vector3<f32>, Vector3<f32>, f32, f32) {
1832        let mut position = Vector3::zeros();
1833        let mut velocity = Vector3::zeros();
1834
1835        let pos_basis = self.basis(t, 0);
1836        let vel_basis = self.basis(t, 1);
1837
1838        let row_offset = self.polyorder * segment;
1839
1840        let coeff_segment = &self.coeff.view((row_offset, 0), (self.polyorder, 4));
1841
1842        let pose_4d = pos_basis.transpose() * coeff_segment;
1843        let vel_4d = vel_basis.transpose() * coeff_segment;
1844
1845        //Convert to f32. This is bad for performance but currently no other way as OSQP requires f64 and Peng utilizes f32.
1846        let pose_4d_f32 = pose_4d.map(|x| x as f32);
1847        let vel_4d_f32 = vel_4d.map(|x| x as f32);
1848
1849        // Extract the necessary values
1850        for i in 0..3 {
1851            position[i] = pose_4d_f32[(0, i)];
1852            velocity[i] = vel_4d_f32[(0, i)];
1853        }
1854        let yaw = pose_4d_f32[(0, 3)];
1855        let yaw_rate = vel_4d_f32[(0, 3)];
1856
1857        (position, velocity, yaw, yaw_rate)
1858    }
1859
1860    /// Solve the Quadratic Programming problem and compute the coefficients for the whole trajectory.
1861    /// The coefficients will be arranged as (poly_order*num_segments, num_dims)
1862    /// # Errors
1863    /// * `SimulationError::OSQPError` - Failed to set up or solve OSQP problem
1864    fn compute_coefficients(&mut self) -> Result<(), SimulationError> {
1865        let settings: Settings = Settings::default()
1866            .max_iter(10000000)
1867            .eps_abs(1e-6)
1868            .eps_rel(5e-6)
1869            .verbose(false);
1870
1871        let (p, q, a, l, u) = self.prepare_qp_problem();
1872
1873        // Create an OSQP problem
1874        let mut problem = Problem::new(p, q.as_slice(), a, l.as_slice(), u.as_slice(), &settings)
1875            .map_err(|_| {
1876            SimulationError::OSQPError("Failed to set up OSQP problem".to_string())
1877        })?;
1878
1879        // Solve problem
1880        let result = problem.solve();
1881
1882        // Obtain the solution coefficients
1883        let coeff_vector = result
1884            .x()
1885            .map(|solution| DVector::from_iterator(solution.len(), solution.iter().cloned()))
1886            .ok_or(SimulationError::OSQPError(
1887                "Unable to solve OSQP problem".to_string(),
1888            ))?;
1889
1890        self.coeff = DMatrix::from_column_slice(
1891            self.segment_times.len() * self.polyorder,
1892            4,
1893            coeff_vector.as_slice(),
1894        );
1895
1896        Ok(())
1897    }
1898
1899    /// Combine the objective function, the equality and inequality constraints into a format for OSQP
1900    /// # Returns
1901    /// * A tuple containing the Q matrix, the linear term vector, the A matrix, the lower bound vector, and the upper bound vector
1902    fn prepare_qp_problem(
1903        &self,
1904    ) -> (
1905        CscMatrix,
1906        DVector<f64>,
1907        CscMatrix,
1908        DVector<f64>,
1909        DVector<f64>,
1910    ) {
1911        // 1. Generate the Q Matrix for x, y, z, yaw. Since the objective function is the same for all, we simply clone the first matrix instead of recomputing for each dim.
1912        let num_dims = 4;
1913        let num_vars_per_dim = self.polyorder * self.segment_times.len();
1914        let num_vars = num_vars_per_dim * num_dims;
1915
1916        // Stack Q matrices into a block-diagonal matrix (joint optimization)
1917        let q_x = self.generate_obj_function();
1918        let mut q = DMatrix::zeros(num_vars, num_vars);
1919
1920        for dim in 0..num_dims {
1921            let row_offset = dim * num_vars_per_dim;
1922            let col_offset = dim * num_vars_per_dim;
1923            q.view_mut((row_offset, col_offset), (q_x.nrows(), q_x.ncols()))
1924                .copy_from(&q_x);
1925        }
1926
1927        // 2. Generate the Equality constraints for x, y, z, yaw
1928
1929        // Construct a block diagonal A_eq matrix
1930        let num_constraints_per_block = (2 + self.smooth_upto) * (self.segment_times.len() + 1) - 2;
1931        let mut a_eq = DMatrix::zeros(num_constraints_per_block * num_dims, num_vars);
1932        let mut b_eq = DVector::zeros(num_constraints_per_block * num_dims);
1933
1934        for dim in 0..num_dims {
1935            let row_offset = dim * num_constraints_per_block;
1936            let col_offset = dim * num_vars_per_dim;
1937            let (small_a_eq, small_b_eq) = self.generate_equality_constraints(dim);
1938            a_eq.view_mut(
1939                (row_offset, col_offset),
1940                (num_constraints_per_block, num_vars_per_dim),
1941            )
1942            .copy_from(&small_a_eq);
1943            b_eq.rows_mut(row_offset, num_constraints_per_block)
1944                .copy_from(&small_b_eq);
1945        }
1946
1947        // 3. Generate the Inequality constraints, similar to how we generate the Q Matrix.
1948        // If no inequality constraints, create empty matrices
1949        let mut a_ineq: DMatrix<f64>;
1950        let mut d: DVector<f64>;
1951        let mut f: DVector<f64>;
1952        if self.max_acceleration > 0.0 && self.max_velocity > 0.0 {
1953            let (a_ineq_x, d_x, f_x) = self.generate_inequality_constraints();
1954            a_ineq = DMatrix::zeros(a_ineq_x.nrows() * num_dims, num_vars);
1955            d = DVector::zeros(d_x.len() * num_dims);
1956            f = DVector::zeros(f_x.len() * num_dims);
1957
1958            for dim in 0..num_dims {
1959                let row_offset = dim * a_ineq_x.nrows();
1960                let col_offset = dim * num_vars_per_dim;
1961                a_ineq
1962                    .view_mut(
1963                        (row_offset, col_offset),
1964                        (a_ineq_x.nrows(), a_ineq_x.ncols()),
1965                    )
1966                    .copy_from(&a_ineq_x);
1967                d.rows_mut(row_offset, d_x.len()).copy_from(&d_x);
1968                f.rows_mut(row_offset, f_x.len()).copy_from(&f_x);
1969            }
1970        } else {
1971            // If no inequality constraints, create empty matrices
1972            a_ineq = DMatrix::zeros(0, num_vars);
1973            d = DVector::zeros(0);
1974            f = DVector::zeros(0);
1975        }
1976
1977        // 4. Combine equality and optionally, inequality constraints
1978        let total_constraints = a_eq.nrows() + a_ineq.nrows();
1979        let mut a = DMatrix::zeros(total_constraints, num_vars);
1980        a.view_mut((0, 0), (a_eq.nrows(), a_eq.ncols()))
1981            .copy_from(&a_eq);
1982
1983        let mut l = DVector::zeros(total_constraints);
1984        let mut u = DVector::zeros(total_constraints);
1985
1986        // Lower and upper bounds
1987        l.rows_mut(0, b_eq.len()).copy_from(&b_eq);
1988        u.rows_mut(0, b_eq.len()).copy_from(&b_eq); // Equality constraints: Enforcing l = u = b_eq is the same as A_eq*x = b_eq
1989
1990        if a_ineq.nrows() > 0 {
1991            a.view_mut((a_eq.nrows(), 0), (a_ineq.nrows(), a_ineq.ncols()))
1992                .copy_from(&a_ineq);
1993            l.rows_mut(b_eq.len(), d.len()).copy_from(&d);
1994            u.rows_mut(b_eq.len(), f.len()).copy_from(&f);
1995        }
1996
1997        // 5. Convert matrices to CSC format for osqp.rs
1998        let q_csc = self.convert_dense_to_sparse(&q).into_upper_tri();
1999        let a_csc = self.convert_dense_to_sparse(&a);
2000
2001        (q_csc, DVector::zeros(num_vars), a_csc, l, u)
2002    }
2003
2004    /// Generate the Q Cost Matrix and A Matrix in the required Csc sparse matrix form for osqp by converting the dense DMatrix.
2005    /// # Arguments
2006    /// * `a` - The dense matrix to be converted.
2007    /// # Returns
2008    /// * The sparse matrix in CSC format.
2009    fn convert_dense_to_sparse(&self, a: &DMatrix<f64>) -> CscMatrix {
2010        let (rows, cols) = a.shape();
2011
2012        let column_major_iter: Vec<f64> = a
2013            .column_iter()
2014            .flat_map(|col| col.iter().copied().collect::<Vec<f64>>())
2015            .collect();
2016
2017        CscMatrix::from_column_iter(rows, cols, column_major_iter)
2018    }
2019
2020    /// Generate the entire objective function for the entire trajectory for a single dimension
2021    /// # Returns
2022    /// * The objective function matrix.
2023    fn generate_obj_function(&self) -> DMatrix<f64> {
2024        let num_segments = self.segment_times.len();
2025        let coeff_num = self.polyorder * num_segments;
2026
2027        let mut q_total: DMatrix<f64> = DMatrix::zeros(coeff_num, coeff_num);
2028
2029        for (i, &segment_time) in self.segment_times.iter().enumerate() {
2030            let q_segment = self.generate_q(self.min_deriv, segment_time);
2031            q_total
2032                .view_mut(
2033                    (i * self.polyorder, i * self.polyorder),
2034                    (self.polyorder, self.polyorder),
2035                )
2036                .copy_from(&q_segment);
2037        }
2038        q_total
2039    }
2040
2041    /// Generate the Q matrix for a single segment of a single dimension
2042    /// # Arguments
2043    /// * `min_deriv` - The minimum derivative order to consider.
2044    /// * `time` - The time at which to evaluate the basis functions.
2045    /// # Returns
2046    /// * The Q matrix for the segment.
2047    fn generate_q(&self, min_deriv: usize, time: f32) -> DMatrix<f64> {
2048        let mut q: DMatrix<f64> = DMatrix::zeros(self.polyorder, self.polyorder);
2049        let poly = self.basis(time, min_deriv);
2050
2051        for i in 0..self.polyorder {
2052            for j in 0..self.polyorder {
2053                let power_order = ((i + j) as f64 - (2 * min_deriv) as f64 + 1.0).max(1.0); // Avoid division by zero
2054                q[(i, j)] = poly[i] * poly[j] * time as f64 / power_order;
2055            }
2056        }
2057        q
2058    }
2059
2060    /// Generate the equality matrix for a single dimension
2061    /// # Arguments
2062    /// * `dimension` - The dimension for which to generate the equality matrix.
2063    /// # Returns
2064    /// * The equality matrix and the vector of constraints.
2065    fn generate_equality_constraints(&self, dimension: usize) -> (DMatrix<f64>, DVector<f64>) {
2066        let num_segments = self.segment_times.len();
2067        let num_coeff = self.polyorder * num_segments;
2068
2069        // 1 position constraint + 4 derivative constraints (upto snap) for each beginning and ending waypoint = 5*2 +
2070        // 2 position constraints + 4 derivative constraints (upto snap) for each intermediate waypoint = 6*num__intermediate_waypoints. This is basically 6*(num_segments+1)-2.
2071        let num_constraints = (2 + self.smooth_upto) * (num_segments + 1) - 2;
2072
2073        // Initialize the equality matrix a and the vector b. a*x = b, where x is the vector of coefficients we want to compute.
2074
2075        let mut a: DMatrix<f64> = DMatrix::zeros(num_constraints, num_coeff);
2076        let mut b: DVector<f64> = DVector::zeros(num_constraints);
2077
2078        let mut row_index: usize = 0;
2079
2080        for i in 0..=num_segments {
2081            let col_offset = i * self.polyorder;
2082
2083            match i {
2084                // First part of the first waypoint of the trajectory. NOTE: The first position is assumed to be the position of the robot. Handles position[0] + derivatives = 0
2085                // This results in 5 constraints (rows) added here.
2086                0 => {
2087                    let row = self.basis(0.0, 0);
2088                    a.row_mut(row_index)
2089                        .columns_mut(col_offset, self.polyorder)
2090                        .copy_from(&row.transpose());
2091                    b[row_index] = self.waypoints[i][dimension] as f64;
2092                    row_index += 1;
2093
2094                    //All derivatives from velocity up until snap must be 0
2095                    for j in 1..=self.smooth_upto {
2096                        let row = self.basis(0.0, j);
2097                        a.row_mut(row_index)
2098                            .columns_mut(col_offset, self.polyorder)
2099                            .copy_from(&row.transpose());
2100                        b[row_index] = 0.0;
2101                        row_index += 1;
2102                    }
2103                }
2104                //Last part of the last waypoint of the trajectory. Handles position[n-1] + derivatives = 0. This results in 5 constraints (rows) added here.
2105                n if n == num_segments => {
2106                    let time = self.segment_times[num_segments - 1];
2107                    let row = self.basis(time, 0);
2108                    a.row_mut(row_index)
2109                        .columns_mut(col_offset - self.polyorder, self.polyorder)
2110                        .copy_from(&row.transpose());
2111                    b[row_index] = self.waypoints[i][dimension] as f64;
2112                    row_index += 1;
2113
2114                    //All derivatives from velocity up until snap must be 0
2115                    for j in 1..=self.smooth_upto {
2116                        let row = self.basis(time, j);
2117                        a.row_mut(row_index)
2118                            .columns_mut(col_offset - self.polyorder, self.polyorder)
2119                            .copy_from(&row.transpose());
2120                        b[row_index] = 0.0;
2121                        row_index += 1;
2122                    }
2123                }
2124                // Intermediate segments. Handles prev_segment_position[i-1] = waypoint[i] = curr_segment_position[i] + derivative[i-1] = derivative[i]
2125                // This results in 6 constraints (rows) added here.
2126                _ => {
2127                    {
2128                        // Position constraint. position at segment i-1 must be equal to position at i which is equal to waypoint[i]
2129                        let segment_prev = self.basis(self.segment_times[i - 1], 0);
2130                        let segment_next = self.basis(0.0, 0);
2131                        let prev_col_offset = (i - 1) * self.polyorder;
2132
2133                        // Offset the col index to i-1 to assign it to the previous segment's coefficients
2134                        a.row_mut(row_index)
2135                            .columns_mut(prev_col_offset, self.polyorder)
2136                            .copy_from(&segment_prev.transpose());
2137                        b[row_index] = self.waypoints[i][dimension] as f64;
2138                        row_index += 1;
2139                        // Offset the col index to i to assign it to this segment's coefficients
2140                        a.row_mut(row_index)
2141                            .columns_mut(col_offset, self.polyorder)
2142                            .copy_from(&segment_next.transpose());
2143                        b[row_index] = self.waypoints[i][dimension] as f64;
2144                        row_index += 1;
2145
2146                        // To ensure velocity upto snap continuity, we equate the derivatives of the previous segment with the derivatives of the next segment.
2147                        // Since we don't care what the value of the derivatives are, we just formulate the constraint such that deriv[prev_segment] - deriv[next_segment] = 0.
2148                        // We do this by placing the derivatives of the prev segment at those times in the same row, along with the negative of next segment's derivatives.
2149
2150                        for j in 1..=self.smooth_upto {
2151                            let segment_prev = self.basis(self.segment_times[i - 1], j);
2152                            let segment_next = self.basis(0.0, j);
2153
2154                            a.row_mut(row_index)
2155                                .columns_mut(prev_col_offset, self.polyorder)
2156                                .copy_from(&segment_prev.transpose());
2157                            a.row_mut(row_index)
2158                                .columns_mut(col_offset, self.polyorder)
2159                                .add_assign(-segment_next.transpose());
2160
2161                            b[row_index] = 0.0;
2162
2163                            row_index += 1;
2164                        }
2165                    }
2166                }
2167            }
2168        }
2169
2170        (a, b)
2171    }
2172    /// Compute the number of constraints
2173    /// # Returns
2174    /// * The total number of constraints.
2175    fn compute_num_constraints(&self) -> usize {
2176        // Skip if no inequality constraints
2177        if self.max_velocity == 0.0 || self.max_acceleration == 0.0 {
2178            return 0;
2179        }
2180        // For each segment, compute number of timesteps
2181        let total_constraints: usize = self
2182            .segment_times
2183            .iter()
2184            .map(|&segment_time| {
2185                // Calculate number of steps, excluding the end point
2186                let num_timesteps = (segment_time / self.dt).floor() as usize;
2187                // 2 constraints (velocity and acceleration) per timestep
2188                num_timesteps * 2
2189            })
2190            .sum();
2191        total_constraints
2192    }
2193
2194    /// Generate inequality constrain for a single dimension. This works by sampling the derivatives using the basis vector across time steps and
2195    /// enforcing velocity and acceleration constraints for each of those basis vectors.
2196    /// # Returns
2197    /// A tuple containing the inequality matrix `c`, the lower bound vector `d`, and the upper bound vector `f`.
2198    fn generate_inequality_constraints(&self) -> (DMatrix<f64>, DVector<f64>, DVector<f64>) {
2199        let num_segments = self.segment_times.len();
2200        let coeff_num = self.polyorder * num_segments;
2201
2202        let num_constraints = self.compute_num_constraints();
2203
2204        // Initialize the inequality matrix c, along with the vectors d and f, representing the inequality d <= c*x <= f, where x is the vector of coefficients.
2205
2206        let mut c: DMatrix<f64> = DMatrix::zeros(num_constraints, coeff_num);
2207        let mut d = DVector::from_element(num_constraints, 0.0); // We fill this up to -max_velocity or -max_acceleration.
2208        let mut f = DVector::from_element(num_constraints, 0.0); // We fill this up.
2209
2210        let mut row_index = 0;
2211
2212        // Iterate over each segment, then each time step in one segment, compute its basis vector for velocity and acceleration, and set the maximum.
2213        for (i, &segment_time) in self.segment_times.iter().enumerate() {
2214            // Col index for each segment
2215            let col_offset = i * self.polyorder;
2216
2217            let mut t = self.dt; // To avoid placing equality and inequality constraints at in the beginning.
2218                                 //Iterate over each time step in one segment
2219            while t < segment_time {
2220                let row_velocity = self.basis(t, 1);
2221
2222                c.row_mut(row_index)
2223                    .columns_mut(col_offset, self.polyorder)
2224                    .copy_from(&row_velocity.transpose());
2225                f[row_index] = self.max_velocity as f64;
2226                d[row_index] = -self.max_velocity as f64;
2227                row_index += 1;
2228
2229                let row_acceleration = self.basis(t, 2);
2230
2231                c.row_mut(row_index)
2232                    .columns_mut(col_offset, self.polyorder)
2233                    .copy_from(&row_acceleration.transpose());
2234                f[row_index] = self.max_acceleration as f64;
2235                d[row_index] = -self.max_acceleration as f64;
2236                row_index += 1;
2237                t += self.dt; // Update time step to compute basis vector at the next step.
2238            }
2239        }
2240        (c, d, f)
2241    }
2242
2243    /// Generate the basis vector of the given derivative order at the given time.
2244    /// # Arguments
2245    /// * `time` - The time at which to evaluate the basis vector.
2246    /// * `derivative` - The derivative order of the basis vector.
2247    /// # Returns
2248    /// * The basis vector of the given derivative order at the given time.
2249    fn basis(&self, time: f32, derivative: usize) -> DVector<f64> {
2250        assert!(derivative <= 4, "Derivative order must be less than 4");
2251        let time_f64 = time as f64;
2252        let t: Vec<f64> = (0..self.polyorder)
2253            .map(|i| time_f64.powi(i as i32))
2254            .collect();
2255
2256        let mut b = DVector::from_element(self.polyorder, 0.0);
2257
2258        for i in 0..self.polyorder {
2259            let power = i as isize - derivative as isize;
2260            let mut coeff = 1.0;
2261
2262            if power < 0 {
2263                b[i] = 0.0;
2264            } else {
2265                for j in (i - derivative + 1..=i).rev() {
2266                    coeff *= j as f64;
2267                }
2268                b[i] = coeff * t[power as usize];
2269            }
2270        }
2271        b
2272    }
2273}
2274
2275/// Implement the `Planner` trait for `QPpolyTrajPlanner`
2276impl Planner for QPpolyTrajPlanner {
2277    fn plan(
2278        &self,
2279        _current_position: Vector3<f32>,
2280        _current_velocity: Vector3<f32>,
2281        time: f32,
2282    ) -> (Vector3<f32>, Vector3<f32>, f32) {
2283        let relative_time = time - self.start_time;
2284        // Find the current segment
2285        let mut segment_start_time = 0.0;
2286        let mut current_segment = 0;
2287        for (i, &segment_duration) in self.segment_times.iter().enumerate() {
2288            if relative_time < segment_start_time + segment_duration {
2289                current_segment = i;
2290                break;
2291            }
2292            segment_start_time += segment_duration;
2293        }
2294        // Evaluate the polynomial for the current segment
2295        let segment_time = relative_time - segment_start_time;
2296        let (position, velocity, yaw, _yaw_rate) =
2297            self.evaluate_polynomial(segment_time, current_segment);
2298        (position, velocity, yaw)
2299    }
2300
2301    fn is_finished(
2302        &self,
2303        current_position: Vector3<f32>,
2304        time: f32,
2305    ) -> Result<bool, SimulationError> {
2306        let last_waypoint = self.waypoints.last().ok_or(SimulationError::OtherError(
2307            "No waypoints available".to_string(),
2308        ))?;
2309
2310        let total_time: f32 = self.segment_times.iter().sum();
2311
2312        let last_position = Vector3::new(last_waypoint[0], last_waypoint[1], last_waypoint[2]);
2313
2314        Ok(time >= self.start_time + total_time && (current_position - last_position).norm() < 0.1)
2315    }
2316}
2317/// Represents a step in the planner schedule.
2318/// # Example
2319/// ```
2320/// use peng_quad::PlannerStepConfig;
2321/// let step = PlannerStepConfig {
2322///     step: 0,
2323///     planner_type: "MinimumJerkLocalPlanner".to_string(),
2324///     params: serde_yaml::Value::Null,
2325/// };
2326/// ```
2327pub struct PlannerStepConfig {
2328    /// The simulation step at which this planner should be activated (in ms unit).
2329    pub step: usize,
2330    /// The type of planner to use for this step.
2331    pub planner_type: String,
2332    /// Additional parameters for the planner, stored as a YAML value.
2333    pub params: serde_yaml::Value,
2334}
2335/// Updates the planner based on the current simulation step and configuration
2336/// # Arguments
2337/// * `planner_manager` - The PlannerManager instance to update
2338/// * `step` - The current simulation step in ms unit
2339/// * `time` - The current simulation time
2340/// * `simulation_frequency' - The simulation frequency in Hz
2341/// * `quad` - The Quadrotor instance
2342/// * `obstacles` - The current obstacles in the simulation
2343/// * `planner_config` - The planner configuration
2344/// # Errors
2345/// * If the planner could not be created
2346/// # Example
2347/// ```
2348/// use peng_quad::{PlannerManager, Quadrotor, Obstacle, PlannerStepConfig, update_planner};
2349/// use nalgebra::Vector3;
2350/// let simulation_frequency = 1000;
2351/// let initial_position = Vector3::new(0.0, 0.0, 0.0);
2352/// let initial_yaw = 0.0;
2353/// let mut planner_manager = PlannerManager::new(initial_position, initial_yaw);
2354/// let step = 0;
2355/// let time = 0.0;
2356/// let (time_step, mass, gravity, drag_coefficient) = (0.01, 1.3, 9.81, 0.01);
2357/// let inertia_matrix = [0.0347563, 0.0, 0.0, 0.0, 0.0458929, 0.0, 0.0, 0.0, 0.0977];
2358/// let quadrotor = Quadrotor::new(time_step, mass, gravity, drag_coefficient, inertia_matrix).unwrap();
2359/// let obstacles = vec![Obstacle::new(Vector3::new(0.0, 0.0, 0.0), Vector3::new(0.0, 0.0, 0.0), 1.0)];
2360/// let planner_config = vec![PlannerStepConfig {
2361///     step: 0,
2362///     planner_type: "MinimumJerkLine".to_string(),
2363///     params:
2364///        serde_yaml::from_str(r#"
2365///        end_position: [0.0, 0.0, 1.0]
2366///        end_yaw: 0.0
2367///        duration: 2.0
2368///        "#).unwrap(),
2369/// }];
2370/// update_planner(&mut planner_manager, step, time, simulation_frequency, &quadrotor, &obstacles, &planner_config).unwrap();
2371/// ```
2372pub fn update_planner(
2373    planner_manager: &mut PlannerManager,
2374    step: usize,
2375    time: f32,
2376    simulation_frequency: usize,
2377    quad: &Quadrotor,
2378    obstacles: &[Obstacle],
2379    planner_config: &[PlannerStepConfig],
2380) -> Result<(), SimulationError> {
2381    if let Some(planner_step) = planner_config
2382        .iter()
2383        .find(|s| s.step * simulation_frequency == step * 1000)
2384    {
2385        log::info!("Time: {:.2} s,\tSwitch {}", time, planner_step.planner_type);
2386        planner_manager.set_planner(create_planner(planner_step, quad, time, obstacles)?);
2387    }
2388    Ok(())
2389}
2390/// Creates a planner based on the configuration
2391/// # Arguments
2392/// * `step` - The configuration for the planner step in ms unit
2393/// * `quad` - The Quadrotor instance
2394/// * `time` - The current simulation time
2395/// * `obstacles` - The current obstacles in the simulation
2396/// # Returns
2397/// * `PlannerType` - The created planner
2398/// # Errors
2399/// * If the planner type is not recognized
2400/// # Example
2401/// ```
2402/// use peng_quad::{PlannerType, Quadrotor, Obstacle, PlannerStepConfig, create_planner};
2403/// use nalgebra::Vector3;
2404/// let step = PlannerStepConfig {
2405///    step: 0,
2406///   planner_type: "MinimumJerkLine".to_string(),
2407///   params:
2408///       serde_yaml::from_str(r#"
2409///       end_position: [0.0, 0.0, 1.0]
2410///       end_yaw: 0.0
2411///       duration: 2.0
2412///       "#).unwrap(),
2413/// };
2414/// let time = 0.0;
2415/// let (time_step, mass, gravity, drag_coefficient) = (0.01, 1.3, 9.81, 0.01);
2416/// let inertia_matrix = [0.0347563, 0.0, 0.0, 0.0, 0.0458929, 0.0, 0.0, 0.0, 0.0977];
2417/// let quadrotor = Quadrotor::new(time_step, mass, gravity, drag_coefficient, inertia_matrix).unwrap();
2418/// let obstacles = vec![Obstacle::new(Vector3::new(0.0, 0.0, 0.0), Vector3::new(0.0, 0.0, 0.0), 1.0)];
2419/// let planner = create_planner(&step, &quadrotor, time, &obstacles).unwrap();
2420/// match planner {
2421///    PlannerType::MinimumJerkLine(_) => log::info!("Created MinimumJerkLine planner"),
2422///   _ => log::info!("Created another planner"),
2423/// }
2424/// ```
2425pub fn create_planner(
2426    step: &PlannerStepConfig,
2427    quad: &Quadrotor,
2428    time: f32,
2429    obstacles: &[Obstacle],
2430) -> Result<PlannerType, SimulationError> {
2431    let params = &step.params;
2432    match step.planner_type.as_str() {
2433        "MinimumJerkLine" => Ok(PlannerType::MinimumJerkLine(MinimumJerkLinePlanner {
2434            start_position: quad.position,
2435            end_position: parse_vector3(params, "end_position")?,
2436            start_yaw: quad.orientation.euler_angles().2,
2437            end_yaw: parse_f32(params, "end_yaw")?,
2438            start_time: time,
2439            duration: parse_f32(params, "duration")?,
2440        })),
2441        "Lissajous" => Ok(PlannerType::Lissajous(LissajousPlanner {
2442            start_position: quad.position,
2443            center: parse_vector3(params, "center")?,
2444            amplitude: parse_vector3(params, "amplitude")?,
2445            frequency: parse_vector3(params, "frequency")?,
2446            phase: parse_vector3(params, "phase")?,
2447            start_time: time,
2448            duration: parse_f32(params, "duration")?,
2449            start_yaw: quad.orientation.euler_angles().2,
2450            end_yaw: parse_f32(params, "end_yaw")?,
2451            ramp_time: parse_f32(params, "ramp_time")?,
2452        })),
2453        "Circle" => Ok(PlannerType::Circle(CirclePlanner {
2454            center: parse_vector3(params, "center")?,
2455            radius: parse_f32(params, "radius")?,
2456            angular_velocity: parse_f32(params, "angular_velocity")?,
2457            start_position: quad.position,
2458            start_time: time,
2459            duration: parse_f32(params, "duration")?,
2460            start_yaw: quad.orientation.euler_angles().2,
2461            end_yaw: quad.orientation.euler_angles().2,
2462            ramp_time: parse_f32(params, "ramp_time")?,
2463        })),
2464        "ObstacleAvoidance" => Ok(PlannerType::ObstacleAvoidance(ObstacleAvoidancePlanner {
2465            target_position: parse_vector3(params, "target_position")?,
2466            start_time: time,
2467            duration: parse_f32(params, "duration")?,
2468            start_yaw: quad.orientation.euler_angles().2,
2469            end_yaw: parse_f32(params, "end_yaw")?,
2470            obstacles: obstacles.to_owned(),
2471            k_att: parse_f32(params, "k_att")?,
2472            k_rep: parse_f32(params, "k_rep")?,
2473            k_vortex: parse_f32(params, "k_vortex")?,
2474            d0: parse_f32(params, "d0")?,
2475            d_target: parse_f32(params, "d_target")?,
2476            max_speed: parse_f32(params, "max_speed")?,
2477        })),
2478        "MinimumSnapWaypoint" => {
2479            let mut waypoints = vec![quad.position];
2480            waypoints.extend(
2481                params["waypoints"]
2482                    .as_sequence()
2483                    .ok_or_else(|| SimulationError::OtherError("Invalid waypoints".to_string()))?
2484                    .iter()
2485                    .map(|w| {
2486                        w.as_sequence()
2487                            .and_then(|coords| {
2488                                Some(Vector3::new(
2489                                    coords[0].as_f64()? as f32,
2490                                    coords[1].as_f64()? as f32,
2491                                    coords[2].as_f64()? as f32,
2492                                ))
2493                            })
2494                            .ok_or(SimulationError::OtherError("Invalid waypoint".to_string()))
2495                    })
2496                    .collect::<Result<Vec<Vector3<f32>>, SimulationError>>()?,
2497            );
2498            let mut yaws = vec![quad.orientation.euler_angles().2];
2499            yaws.extend(
2500                params["yaws"]
2501                    .as_sequence()
2502                    .ok_or(SimulationError::OtherError("Invalid yaws".to_string()))?
2503                    .iter()
2504                    .map(|y| {
2505                        y.as_f64()
2506                            .map(|v| v as f32)
2507                            .ok_or(SimulationError::OtherError("Invalid yaw".to_string()))
2508                    })
2509                    .collect::<Result<Vec<f32>, SimulationError>>()?,
2510            );
2511            let segment_times = params["segment_times"]
2512                .as_sequence()
2513                .ok_or_else(|| SimulationError::OtherError("Invalid segment_times".to_string()))?
2514                .iter()
2515                .map(|t| {
2516                    t.as_f64().map(|v| v as f32).ok_or_else(|| {
2517                        SimulationError::OtherError("Invalid segment time".to_string())
2518                    })
2519                })
2520                .collect::<Result<Vec<f32>, SimulationError>>()?;
2521            MinimumSnapWaypointPlanner::new(waypoints, yaws, segment_times, time)
2522                .map(PlannerType::MinimumSnapWaypoint)
2523        }
2524        "QPpolyTraj" => {
2525            let mut waypoints = vec![vec![
2526                quad.position.x,
2527                quad.position.y,
2528                quad.position.z,
2529                quad.orientation.euler_angles().2,
2530            ]];
2531            waypoints.extend(
2532                params["waypoints"]
2533                    .as_sequence()
2534                    .ok_or_else(|| SimulationError::OtherError("Invalid waypoints".to_string()))?
2535                    .iter()
2536                    .map(|w| {
2537                        w.as_sequence()
2538                            .and_then(|coords| {
2539                                Some(vec![
2540                                    coords[0].as_f64()? as f32,
2541                                    coords[1].as_f64()? as f32,
2542                                    coords[2].as_f64()? as f32,
2543                                    coords[3].as_f64()? as f32,
2544                                ])
2545                            })
2546                            .ok_or(SimulationError::OtherError("Invalid waypoint".to_string()))
2547                    })
2548                    .collect::<Result<Vec<Vec<f32>>, SimulationError>>()?,
2549            );
2550
2551            let segment_times = params["segment_times"]
2552                .as_sequence()
2553                .ok_or_else(|| SimulationError::OtherError("Invalid segment_times".to_string()))?
2554                .iter()
2555                .map(|t| {
2556                    t.as_f64().map(|v| v as f32).ok_or_else(|| {
2557                        SimulationError::OtherError("Invalid segment time".to_string())
2558                    })
2559                })
2560                .collect::<Result<Vec<f32>, SimulationError>>()?;
2561            QPpolyTrajPlanner::new(
2562                waypoints,
2563                segment_times,
2564                parse_usize(params, "polynomial_order")?,
2565                parse_usize(params, "minimize_derivative")?,
2566                parse_usize(params, "smooth_upto")?,
2567                parse_f32(params, "max_velocity")?,
2568                parse_f32(params, "max_acceleration")?,
2569                time,
2570                parse_f32(params, "dt")?,
2571            )
2572            .map(PlannerType::QPpolyTraj)
2573        }
2574        "Landing" => Ok(PlannerType::Landing(LandingPlanner {
2575            start_position: quad.position,
2576            start_time: time,
2577            duration: parse_f32(params, "duration")?,
2578            start_yaw: quad.orientation.euler_angles().2,
2579        })),
2580        _ => Err(SimulationError::OtherError(format!(
2581            "Unknown planner type: {}",
2582            step.planner_type
2583        ))),
2584    }
2585}
2586/// Helper function to parse Vector3 from YAML
2587/// # Arguments
2588/// * `value` - YAML value
2589/// * `key` - key to parse
2590/// # Returns
2591/// * `Vector3<f32>` - parsed vector
2592/// # Errors
2593/// * `SimulationError` - if the value is not a valid vector
2594/// # Example
2595/// ```
2596/// use nalgebra::Vector3;
2597/// use peng_quad::{parse_vector3, SimulationError};
2598/// let value = serde_yaml::from_str("test: [1.0, 2.0, 3.0]").unwrap();
2599/// assert_eq!(parse_vector3(&value, "test").unwrap(), Vector3::new(1.0, 2.0, 3.0));
2600/// ```
2601pub fn parse_vector3(
2602    value: &serde_yaml::Value,
2603    key: &str,
2604) -> Result<Vector3<f32>, SimulationError> {
2605    value[key]
2606        .as_sequence()
2607        .and_then(|seq| {
2608            if seq.len() == 3 {
2609                Some(Vector3::new(
2610                    seq[0].as_f64()? as f32,
2611                    seq[1].as_f64()? as f32,
2612                    seq[2].as_f64()? as f32,
2613                ))
2614            } else {
2615                None
2616            }
2617        })
2618        .ok_or_else(|| SimulationError::OtherError(format!("Invalid {key} vector")))
2619}
2620/// Helper function to parse f32 from YAML
2621/// # Arguments
2622/// * `value` - YAML value
2623/// * `key` - key to parse
2624/// # Returns
2625/// * `f32` - parsed value
2626/// # Errors
2627/// * `SimulationError` - if the value is not a valid f32
2628/// # Example
2629/// ```
2630/// use peng_quad::{parse_f32, SimulationError};
2631/// let value = serde_yaml::from_str("key: 1.0").unwrap();
2632/// let result = parse_f32(&value, "key").unwrap();
2633/// assert_eq!(result, 1.0);
2634/// ```
2635pub fn parse_f32(value: &serde_yaml::Value, key: &str) -> Result<f32, SimulationError> {
2636    value[key]
2637        .as_f64()
2638        .map(|v| v as f32)
2639        .ok_or_else(|| SimulationError::OtherError(format!("Invalid {key}")))
2640}
2641
2642/// Helper function to parse unsigned integer from YAML
2643/// # Arguments
2644/// * `value` - YAML value
2645/// * `key` - key to parse
2646/// # Returns
2647/// * `usize` - parsed value
2648/// # Errors
2649/// * `SimulationError` - if the value is not a valid unsigned integer
2650/// # Example
2651/// ```
2652/// use peng_quad::{parse_usize, SimulationError};
2653/// let value = serde_yaml::from_str("key: 1").unwrap();
2654/// let result = parse_usize(&value, "key").unwrap();
2655/// assert_eq!(result, 1);
2656/// ```
2657pub fn parse_usize(value: &serde_yaml::Value, key: &str) -> Result<usize, SimulationError> {
2658    value[key]
2659        .as_i64()
2660        .and_then(|v| if v >= 0 { Some(v as usize) } else { None }) // Ensure non-negative
2661        .ok_or_else(|| SimulationError::OtherError(format!("Invalid {key}")))
2662}
2663/// Represents an obstacle in the simulation
2664/// # Example
2665/// ```
2666/// use peng_quad::Obstacle;
2667/// use nalgebra::Vector3;
2668/// let obstacle = Obstacle::new(Vector3::new(0.0, 0.0, 0.0), Vector3::new(0.0, 0.0, 0.0), 1.0);
2669/// ```
2670#[derive(Clone)]
2671pub struct Obstacle {
2672    /// The position of the obstacle
2673    pub position: Vector3<f32>,
2674    /// The velocity of the obstacle
2675    pub velocity: Vector3<f32>,
2676    /// The radius of the obstacle
2677    pub radius: f32,
2678}
2679/// Implementation of the Obstacle
2680impl Obstacle {
2681    /// Creates a new obstacle with the given position, velocity, and radius
2682    /// # Arguments
2683    /// * `position` - The position of the obstacle
2684    /// * `velocity` - The velocity of the obstacle
2685    /// * `radius` - The radius of the obstacle
2686    /// # Returns
2687    /// * The new obstacle instance
2688    /// # Example
2689    /// ```
2690    /// use peng_quad::Obstacle;
2691    /// use nalgebra::Vector3;
2692    /// let obstacle = Obstacle::new(Vector3::new(0.0, 0.0, 0.0), Vector3::new(0.0, 0.0, 0.0), 1.0);
2693    /// ```
2694    pub fn new(position: Vector3<f32>, velocity: Vector3<f32>, radius: f32) -> Self {
2695        Self {
2696            position,
2697            velocity,
2698            radius,
2699        }
2700    }
2701}
2702/// Represents a maze in the simulation
2703/// # Example
2704/// ```
2705/// use peng_quad::{Maze, Obstacle};
2706/// use rand_chacha::ChaCha8Rng;
2707/// use rand::SeedableRng;
2708/// use nalgebra::Vector3;
2709/// let maze = Maze {
2710///     lower_bounds: [0.0, 0.0, 0.0],
2711///     upper_bounds: [1.0, 1.0, 1.0],
2712///     obstacles: vec![Obstacle::new(Vector3::new(0.0, 0.0, 0.0), Vector3::new(0.0, 0.0, 0.0), 1.0)],
2713///     obstacles_velocity_bounds: [0.0, 0.0, 0.0],
2714///     obstacles_radius_bounds: [0.0, 0.0],
2715///     rng: ChaCha8Rng::from_os_rng(),
2716/// };
2717/// ```
2718pub struct Maze {
2719    /// The lower bounds of the maze in the x, y, and z directions
2720    pub lower_bounds: [f32; 3],
2721    /// The upper bounds of the maze in the x, y, and z directions
2722    pub upper_bounds: [f32; 3],
2723    /// The obstacles in the maze
2724    pub obstacles: Vec<Obstacle>,
2725    /// The bounds of the obstacles' velocity
2726    pub obstacles_velocity_bounds: [f32; 3],
2727    /// The bounds of the obstacles' radius
2728    pub obstacles_radius_bounds: [f32; 2],
2729    /// Rng for generating random numbers
2730    pub rng: ChaCha8Rng,
2731}
2732/// Implementation of the maze
2733impl Maze {
2734    /// Creates a new maze with the given bounds and number of obstacles
2735    /// # Arguments
2736    /// * `lower_bounds` - The lower bounds of the maze
2737    /// * `upper_bounds` - The upper bounds of the maze
2738    /// * `num_obstacles` - The number of obstacles in the maze
2739    /// # Returns
2740    /// * The new maze instance
2741    /// # Example
2742    /// ```
2743    /// use peng_quad::Maze;
2744    /// let maze = Maze::new([-1.0, -1.0, -1.0], [1.0, 1.0, 1.0], 5, [0.1, 0.1, 0.1], [0.1, 0.5]);
2745    /// ```
2746    pub fn new(
2747        lower_bounds: [f32; 3],
2748        upper_bounds: [f32; 3],
2749        num_obstacles: usize,
2750        obstacles_velocity_bounds: [f32; 3],
2751        obstacles_radius_bounds: [f32; 2],
2752    ) -> Self {
2753        let mut maze = Maze {
2754            lower_bounds,
2755            upper_bounds,
2756            obstacles: Vec::new(),
2757            obstacles_velocity_bounds,
2758            obstacles_radius_bounds,
2759            rng: ChaCha8Rng::from_os_rng(),
2760        };
2761        maze.generate_obstacles(num_obstacles);
2762        maze
2763    }
2764    /// Generates the obstacles in the maze
2765    /// # Arguments
2766    /// * `num_obstacles` - The number of obstacles to generate
2767    /// # Example
2768    /// ```
2769    /// use peng_quad::Maze;
2770    /// let mut maze = Maze::new([-1.0, -1.0, -1.0], [1.0, 1.0, 1.0], 5, [0.1, 0.1, 0.1], [0.1, 0.5]);
2771    /// maze.generate_obstacles(5);
2772    /// ```
2773    pub fn generate_obstacles(&mut self, num_obstacles: usize) {
2774        self.obstacles = (0..num_obstacles)
2775            .map(|_| {
2776                let position = Vector3::new(
2777                    rand::Rng::random_range(
2778                        &mut self.rng,
2779                        self.lower_bounds[0]..self.upper_bounds[0],
2780                    ),
2781                    rand::Rng::random_range(
2782                        &mut self.rng,
2783                        self.lower_bounds[1]..self.upper_bounds[1],
2784                    ),
2785                    rand::Rng::random_range(
2786                        &mut self.rng,
2787                        self.lower_bounds[2]..self.upper_bounds[2],
2788                    ),
2789                );
2790                let v_bounds = self.obstacles_velocity_bounds;
2791                let r_bounds = self.obstacles_radius_bounds;
2792                let velocity = Vector3::new(
2793                    rand::Rng::random_range(&mut self.rng, -v_bounds[0]..v_bounds[0]),
2794                    rand::Rng::random_range(&mut self.rng, -v_bounds[1]..v_bounds[1]),
2795                    rand::Rng::random_range(&mut self.rng, -v_bounds[2]..v_bounds[2]),
2796                );
2797                let radius = rand::Rng::random_range(&mut self.rng, r_bounds[0]..r_bounds[1]);
2798                Obstacle::new(position, velocity, radius)
2799            })
2800            .collect();
2801    }
2802    /// Updates the obstacles in the maze, if an obstacle hits a boundary, it bounces off
2803    /// # Arguments
2804    /// * `dt` - The time step
2805    /// # Example
2806    /// ```
2807    /// use peng_quad::Maze;
2808    /// let mut maze = Maze::new([-1.0, -1.0, -1.0], [1.0, 1.0, 1.0], 5, [0.1, 0.1, 0.1], [0.1, 0.5]);
2809    /// maze.update_obstacles(0.1);
2810    /// ```
2811    pub fn update_obstacles(&mut self, dt: f32) {
2812        self.obstacles.iter_mut().for_each(|obstacle| {
2813            obstacle.position += obstacle.velocity * dt;
2814            for i in 0..3 {
2815                if obstacle.position[i] - obstacle.radius < self.lower_bounds[i]
2816                    || obstacle.position[i] + obstacle.radius > self.upper_bounds[i]
2817                {
2818                    obstacle.velocity[i] *= -1.0;
2819                }
2820            }
2821        });
2822    }
2823}
2824/// Represents a camera in the simulation which is used to render the depth of the scene
2825/// # Example
2826/// ```
2827/// use peng_quad::Camera;
2828/// let camera = Camera::new((800, 600), 60.0, 0.1, 100.0);
2829/// ```
2830pub struct Camera {
2831    /// The resolution of the camera
2832    pub resolution: (usize, usize),
2833    /// The vertical field of view of the camera
2834    pub fov_vertical: f32,
2835    /// The horizontal field of view of the camera
2836    pub fov_horizontal: f32,
2837    /// The vertical focal length of the camera
2838    pub vertical_focal_length: f32,
2839    /// The horizontal focal length of the camera
2840    pub horizontal_focal_length: f32,
2841    /// The near clipping plane of the camera
2842    pub near: f32,
2843    /// The far clipping plane of the camera
2844    pub far: f32,
2845    /// The aspect ratio of the camera
2846    pub aspect_ratio: f32,
2847    /// The ray directions of each pixel in the camera
2848    pub ray_directions: Vec<Vector3<f32>>,
2849    /// Depth buffer
2850    pub depth_buffer: Vec<f32>,
2851}
2852/// Implementation of the camera
2853impl Camera {
2854    /// Creates a new camera with the given resolution, field of view, near and far clipping planes
2855    /// # Arguments
2856    /// * `resolution` - The resolution of the camera
2857    /// * `fov_vertical` - The vertical field of view of the camera
2858    /// * `near` - The near clipping plane of the camera
2859    /// * `far` - The far clipping plane of the camera
2860    /// # Returns
2861    /// * The new camera instance
2862    /// # Example
2863    /// ```
2864    /// use peng_quad::Camera;
2865    /// let camera = Camera::new((800, 600), 1.0, 5.0, 120.0);
2866    /// ```
2867    pub fn new(resolution: (usize, usize), fov_vertical: f32, near: f32, far: f32) -> Self {
2868        let (width, height) = resolution;
2869        let (aspect_ratio, tan_half_fov) =
2870            (width as f32 / height as f32, (fov_vertical / 2.0).tan());
2871        let mut ray_directions = Vec::with_capacity(width * height);
2872        for y in 0..height {
2873            for x in 0..width {
2874                let x_ndc = (2.0 * x as f32 / width as f32 - 1.0) * aspect_ratio * tan_half_fov;
2875                let y_ndc = (1.0 - 2.0 * y as f32 / height as f32) * tan_half_fov;
2876                ray_directions.push(Vector3::new(1.0, x_ndc, y_ndc).normalize());
2877            }
2878        }
2879        let fov_horizontal =
2880            (width as f32 / height as f32 * (fov_vertical / 2.0).tan()).atan() * 2.0;
2881        let horizontal_focal_length = (width as f32 / 2.0) / ((fov_horizontal / 2.0).tan());
2882        let vertical_focal_length = (height as f32 / 2.0) / ((fov_vertical / 2.0).tan());
2883        let depth_buffer = vec![0.0; width * height];
2884
2885        Self {
2886            resolution,
2887            fov_vertical,
2888            fov_horizontal,
2889            vertical_focal_length,
2890            horizontal_focal_length,
2891            near,
2892            far,
2893            aspect_ratio,
2894            ray_directions,
2895            depth_buffer,
2896        }
2897    }
2898
2899    /// Renders the depth of the scene from the perspective of the quadrotor
2900    ///
2901    /// When the depth value is out of the near and far clipping planes, it is set to infinity
2902    /// When the resolution is larger than 32x24, multi-threading can accelerate the rendering
2903    /// # Arguments
2904    /// * `quad_position` - The position of the quadrotor
2905    /// * `quad_orientation` - The orientation of the quadrotor
2906    /// * `maze` - The maze in the scene
2907    /// * `use_multi_threading` - Whether to use multi-threading to render the depth
2908    /// # Errors
2909    /// * If the depth buffer is not large enough to store the depth values
2910    /// # Example
2911    /// ```
2912    /// use peng_quad::{Camera, Maze};
2913    /// use nalgebra::{Vector3, UnitQuaternion};
2914    /// let mut camera = Camera::new((800, 600), 60.0, 0.1, 100.0);
2915    /// let quad_position = Vector3::new(0.0, 0.0, 0.0);
2916    /// let quad_orientation = UnitQuaternion::identity();
2917    /// let mut maze = Maze::new([-1.0, -1.0, -1.0], [1.0, 1.0, 1.0], 5, [0.1, 0.1, 0.1], [0.1, 0.5]);
2918    /// let use_multi_threading = true;
2919    /// camera.render_depth(&quad_position, &quad_orientation, &maze, use_multi_threading);
2920    /// let use_multi_threading = false;
2921    /// camera.render_depth(&quad_position, &quad_orientation, &maze, use_multi_threading);
2922    /// ```
2923    pub fn render_depth(
2924        &mut self,
2925        quad_position: &Vector3<f32>,
2926        quad_orientation: &UnitQuaternion<f32>,
2927        maze: &Maze,
2928        use_multi_threading: bool,
2929    ) -> Result<(), SimulationError> {
2930        let (width, height) = self.resolution;
2931        let total_pixels = width * height;
2932        let rotation_camera_to_world = quad_orientation.to_rotation_matrix().matrix()
2933            * Matrix3::new(1.0, 0.0, 0.0, 0.0, -1.0, 0.0, 0.0, 0.0, 1.0);
2934        let rotation_world_to_camera = rotation_camera_to_world.transpose();
2935
2936        const CHUNK_SIZE: usize = 64;
2937        if use_multi_threading {
2938            self.depth_buffer
2939                .par_chunks_mut(CHUNK_SIZE)
2940                .enumerate()
2941                .try_for_each(|(chunk_idx, chunk)| {
2942                    let start_idx = chunk_idx * CHUNK_SIZE;
2943                    for (i, depth) in chunk.iter_mut().enumerate() {
2944                        let ray_idx = start_idx + i;
2945                        if ray_idx >= total_pixels {
2946                            break;
2947                        }
2948                        *depth = ray_cast(
2949                            quad_position,
2950                            &rotation_world_to_camera,
2951                            &(rotation_camera_to_world * self.ray_directions[ray_idx]),
2952                            maze,
2953                            self.near,
2954                            self.far,
2955                        )?;
2956                    }
2957                    Ok::<(), SimulationError>(())
2958                })?;
2959        } else {
2960            for i in 0..total_pixels {
2961                self.depth_buffer[i] = ray_cast(
2962                    quad_position,
2963                    &rotation_world_to_camera,
2964                    &(rotation_camera_to_world * self.ray_directions[i]),
2965                    maze,
2966                    self.near,
2967                    self.far,
2968                )?;
2969            }
2970        }
2971        Ok(())
2972    }
2973}
2974
2975/// Casts a ray from the camera origin in the given direction
2976/// # Arguments
2977/// * `origin` - The origin of the ray
2978/// * `rotation_world_to_camera` - The rotation matrix from world to camera coordinates
2979/// * `direction` - The direction of the ray
2980/// * `maze` - The maze in the scene
2981/// * `near` - The minimum distance to consider
2982/// * `far` - The maximum distance to consider
2983/// # Returns
2984/// * The distance to the closest obstacle hit by the ray
2985/// # Errors
2986/// * If the ray does not hit any obstacles
2987/// # Example
2988/// ```
2989/// use peng_quad::{ray_cast, Maze};
2990/// use nalgebra::{Vector3, Matrix3};
2991/// let origin = Vector3::new(0.0, 0.0, 0.0);
2992/// let rotation_world_to_camera = Matrix3::identity();
2993/// let direction = Vector3::new(0.0, 0.0, 1.0);
2994/// let maze = Maze::new([-1.0, -1.0, -1.0], [1.0, 1.0, 1.0], 5, [0.1, 0.1, 0.1], [0.1, 0.5]);
2995/// let near = 0.1;
2996/// let far = 100.0;
2997/// let distance = ray_cast(&origin, &rotation_world_to_camera, &direction, &maze, near, far);
2998/// ```
2999pub fn ray_cast(
3000    origin: &Vector3<f32>,
3001    rotation_world_to_camera: &Matrix3<f32>,
3002    direction: &Vector3<f32>,
3003    maze: &Maze,
3004    near: f32,
3005    far: f32,
3006) -> Result<f32, SimulationError> {
3007    let mut closest_hit = far;
3008    // Inline tube intersection
3009    for axis in 0..3 {
3010        let t_near = if direction[axis] > 0.0 {
3011            (maze.upper_bounds[axis] - origin[axis]) / direction[axis]
3012        } else {
3013            (maze.lower_bounds[axis] - origin[axis]) / direction[axis]
3014        };
3015        if t_near > near && t_near < closest_hit {
3016            let intersection_point = origin + direction * t_near;
3017            let mut valid = true;
3018            for i in 0..3 {
3019                if i != axis
3020                    && (intersection_point[i] < maze.lower_bounds[i]
3021                        || intersection_point[i] > maze.upper_bounds[i])
3022                {
3023                    valid = false;
3024                    break;
3025                }
3026            }
3027            if valid {
3028                closest_hit = t_near;
3029            }
3030        }
3031    }
3032    // Early exit if we've hit a wall closer than any possible obstacle
3033    if closest_hit <= near {
3034        return Ok(f32::INFINITY);
3035    }
3036    // Inline sphere intersection
3037    for obstacle in &maze.obstacles {
3038        let oc = origin - obstacle.position;
3039        let b = oc.dot(direction);
3040        let c = oc.dot(&oc) - obstacle.radius * obstacle.radius;
3041        let discriminant = b * b - c;
3042        if discriminant >= 0.0 {
3043            // let t = -b - discriminant.sqrt();
3044            let t = -b - fast_sqrt(discriminant);
3045            if t > near && t < closest_hit {
3046                closest_hit = t;
3047            }
3048        }
3049    }
3050    if closest_hit < far {
3051        Ok((rotation_world_to_camera * direction * closest_hit).x)
3052    } else {
3053        Ok(f32::INFINITY)
3054    }
3055}
3056/// Logs simulation data to the rerun recording stream
3057/// # Arguments
3058/// * `rec` - The rerun::RecordingStream instance
3059/// * `quad` - The Quadrotor instance
3060/// * `desired_position` - The desired position vector
3061/// * `measured_accel` - The measured acceleration vector
3062/// * `measured_gyro` - The measured angular velocity vector
3063/// # Errors
3064/// * If the data cannot be logged to the recording stream
3065/// # Example
3066/// ```no_run
3067/// use peng_quad::{Quadrotor, log_data};
3068/// use nalgebra::Vector3;
3069/// let rec = rerun::RecordingStreamBuilder::new("peng").spawn().unwrap();
3070/// let (time_step, mass, gravity, drag_coefficient) = (0.01, 1.3, 9.81, 0.01);
3071/// let inertia_matrix = [0.0347563, 0.0, 0.0, 0.0, 0.0458929, 0.0, 0.0, 0.0, 0.0977];
3072/// let quad = Quadrotor::new(time_step, mass, gravity, drag_coefficient, inertia_matrix).unwrap();
3073/// let desired_position = Vector3::new(0.0, 0.0, 0.0);
3074/// let desired_velocity = Vector3::new(0.0, 0.0, 0.0);
3075/// let measured_accel = Vector3::new(0.0, 0.0, 0.0);
3076/// let measured_gyro = Vector3::new(0.0, 0.0, 0.0);
3077/// log_data(&rec, &quad, &desired_position, &desired_velocity, &measured_accel, &measured_gyro).unwrap();
3078/// ```
3079pub fn log_data(
3080    rec: &rerun::RecordingStream,
3081    quad: &Quadrotor,
3082    desired_position: &Vector3<f32>,
3083    desired_velocity: &Vector3<f32>,
3084    measured_accel: &Vector3<f32>,
3085    measured_gyro: &Vector3<f32>,
3086) -> Result<(), SimulationError> {
3087    rec.log(
3088        "world/quad/desired_position",
3089        &rerun::Points3D::new([(desired_position.x, desired_position.y, desired_position.z)])
3090            .with_radii([0.1])
3091            .with_colors([rerun::Color::from_rgb(255, 255, 255)]),
3092    )?;
3093    rec.log(
3094        "world/quad/base_link",
3095        &rerun::Transform3D::from_translation_rotation(
3096            rerun::Vec3D::new(quad.position.x, quad.position.y, quad.position.z),
3097            rerun::Quaternion::from_xyzw([
3098                quad.orientation.i,
3099                quad.orientation.j,
3100                quad.orientation.k,
3101                quad.orientation.w,
3102            ]),
3103        )
3104        .with_axis_length(0.7),
3105    )?;
3106    let (quad_roll, quad_pitch, quad_yaw) = quad.orientation.euler_angles();
3107    let quad_euler_angles: Vector3<f32> = Vector3::new(quad_roll, quad_pitch, quad_yaw);
3108    for (pre, vec) in [
3109        ("position", &quad.position),
3110        ("velocity", &quad.velocity),
3111        ("accel", measured_accel),
3112        ("orientation", &quad_euler_angles),
3113        ("gyro", measured_gyro),
3114        ("desired_position", desired_position),
3115        ("desired_velocity", desired_velocity),
3116    ] {
3117        for (i, a) in ["x", "y", "z"].iter().enumerate() {
3118            rec.log(format!("{pre}/{a}"), &rerun::Scalars::single(vec[i] as f64))?;
3119        }
3120    }
3121    Ok(())
3122}
3123/// Log the maze tube to the rerun recording stream
3124/// # Arguments
3125/// * `rec` - The rerun::RecordingStream instance
3126/// * `maze` - The maze instance
3127/// # Errors
3128/// * If the data cannot be logged to the recording stream
3129/// # Example
3130/// ```no_run
3131/// use peng_quad::{Maze, log_maze_tube};
3132/// use rerun::RecordingStreamBuilder;
3133/// let rec = rerun::RecordingStreamBuilder::new("peng").spawn().unwrap();
3134/// let mut maze = Maze::new([-1.0, -1.0, -1.0], [1.0, 1.0, 1.0], 5, [0.1, 0.1, 0.1], [0.1, 0.5]);
3135/// log_maze_tube(&rec, &maze).unwrap();
3136/// ```
3137pub fn log_maze_tube(rec: &rerun::RecordingStream, maze: &Maze) -> Result<(), SimulationError> {
3138    let (lower_bounds, upper_bounds) = (maze.lower_bounds, maze.upper_bounds);
3139    let center_position = rerun::external::glam::Vec3::new(
3140        (lower_bounds[0] + upper_bounds[0]) / 2.0,
3141        (lower_bounds[1] + upper_bounds[1]) / 2.0,
3142        (lower_bounds[2] + upper_bounds[2]) / 2.0,
3143    );
3144    let half_sizes = rerun::external::glam::Vec3::new(
3145        (upper_bounds[0] - lower_bounds[0]) / 2.0,
3146        (upper_bounds[1] - lower_bounds[1]) / 2.0,
3147        (upper_bounds[2] - lower_bounds[2]) / 2.0,
3148    );
3149    rec.log(
3150        "world/maze/tube",
3151        &rerun::Boxes3D::from_centers_and_half_sizes([center_position], [half_sizes])
3152            .with_colors([rerun::Color::from_rgb(128, 128, 255)]),
3153    )?;
3154    Ok(())
3155}
3156/// Log the maze obstacles to the rerun recording stream
3157/// # Arguments
3158/// * `rec` - The rerun::RecordingStream instance
3159/// * `maze` - The maze instance
3160/// # Errors
3161/// * If the data cannot be logged to the recording stream
3162/// # Example
3163/// ```no_run
3164/// use peng_quad::{Maze, log_maze_obstacles};
3165/// let rec = rerun::RecordingStreamBuilder::new("peng").spawn().unwrap();
3166/// let mut maze = Maze::new([-1.0, -1.0, -1.0], [1.0, 1.0, 1.0], 5, [0.1, 0.1, 0.1], [0.1, 0.5]);
3167/// log_maze_obstacles(&rec, &maze).unwrap();
3168/// ```
3169pub fn log_maze_obstacles(
3170    rec: &rerun::RecordingStream,
3171    maze: &Maze,
3172) -> Result<(), SimulationError> {
3173    let (positions, radii): (Vec<(f32, f32, f32)>, Vec<f32>) = maze
3174        .obstacles
3175        .iter()
3176        .map(|obstacle| {
3177            let pos = obstacle.position;
3178            ((pos.x, pos.y, pos.z), obstacle.radius)
3179        })
3180        .unzip();
3181    rec.log(
3182        "world/maze/obstacles",
3183        &rerun::Points3D::new(positions)
3184            .with_radii(radii)
3185            .with_colors([rerun::Color::from_rgb(255, 128, 128)]),
3186    )?;
3187    Ok(())
3188}
3189/// A struct to hold trajectory data
3190/// # Example
3191/// ```
3192/// use peng_quad::Trajectory;
3193/// let initial_point = nalgebra::Vector3::new(0.0, 0.0, 0.0);
3194/// let mut trajectory = Trajectory::new(initial_point);
3195/// ```
3196pub struct Trajectory {
3197    /// A vector of 3D points
3198    pub points: Vec<Vector3<f32>>,
3199    /// The last point that was logged
3200    pub last_logged_point: Vector3<f32>,
3201    /// The minimum distance between points to log
3202    pub min_distance_threadhold: f32,
3203}
3204/// Implement the Trajectory struct
3205impl Trajectory {
3206    /// Create a new Trajectory instance
3207    /// # Arguments
3208    /// * `initial_point` - The initial point to add to the trajectory
3209    /// # Returns
3210    /// * A new Trajectory instance
3211    /// # Example
3212    /// ```
3213    /// use peng_quad::Trajectory;
3214    /// let initial_point = nalgebra::Vector3::new(0.0, 0.0, 0.0);
3215    /// let mut trajectory = Trajectory::new(initial_point);
3216    /// ```
3217    pub fn new(initial_point: Vector3<f32>) -> Self {
3218        Self {
3219            points: vec![initial_point],
3220            last_logged_point: initial_point,
3221            min_distance_threadhold: 0.05,
3222        }
3223    }
3224    /// Add a point to the trajectory if it is further than the minimum distance threshold
3225    /// # Arguments
3226    /// * `point` - The point to add
3227    /// # Returns
3228    /// * `true` if the point was added, `false` otherwise
3229    /// # Example
3230    /// ```
3231    /// use peng_quad::Trajectory;
3232    /// let mut trajectory = Trajectory::new(nalgebra::Vector3::new(0.0, 0.0, 0.0));
3233    /// let point = nalgebra::Vector3::new(1.0, 0.0, 0.0);
3234    /// assert_eq!(trajectory.add_point(point), true);
3235    /// assert_eq!(trajectory.add_point(point), false);
3236    /// ```
3237    pub fn add_point(&mut self, point: Vector3<f32>) -> bool {
3238        if (point - self.last_logged_point).norm() > self.min_distance_threadhold {
3239            self.points.push(point);
3240            self.last_logged_point = point;
3241            true
3242        } else {
3243            false
3244        }
3245    }
3246}
3247/// log trajectory data to the rerun recording stream
3248/// # Arguments
3249/// * `rec` - The rerun::RecordingStream instance
3250/// * `trajectory` - The Trajectory instance
3251/// # Errors
3252/// * If the data cannot be logged to the recording stream
3253/// # Example
3254/// ```no_run
3255/// use peng_quad::{Trajectory, log_trajectory};
3256/// use nalgebra::Vector3;
3257/// let rec = rerun::RecordingStreamBuilder::new("peng").spawn().unwrap();
3258/// let mut trajectory = Trajectory::new(nalgebra::Vector3::new(0.0, 0.0, 0.0));
3259/// trajectory.add_point(nalgebra::Vector3::new(1.0, 0.0, 0.0));
3260/// log_trajectory(&rec, &trajectory).unwrap();
3261/// ```
3262pub fn log_trajectory(
3263    rec: &rerun::RecordingStream,
3264    trajectory: &Trajectory,
3265) -> Result<(), SimulationError> {
3266    let path = trajectory
3267        .points
3268        .iter()
3269        .map(|p| (p.x, p.y, p.z))
3270        .collect::<Vec<(f32, f32, f32)>>();
3271    rec.log(
3272        "world/quad/path",
3273        &rerun::LineStrips3D::new([path]).with_colors([rerun::Color::from_rgb(0, 255, 255)]),
3274    )?;
3275    Ok(())
3276}
3277/// Log depth image data to the rerun recording stream
3278///
3279/// When the depth value is `f32::INFINITY`, the pixel is considered invalid and logged as black
3280/// When the resolution is larger than 32x24, multi-threading can accelerate the rendering
3281/// # Arguments
3282/// * `rec` - The rerun::RecordingStream instance
3283/// * `cam` - The Camera instance
3284/// * `use_multi_threading` - Whether to use multithreading to log the depth image
3285/// # Errors
3286/// * If the data cannot be logged to the recording stream
3287/// # Example
3288/// ```no_run
3289/// use peng_quad::{log_depth_image, Camera};
3290/// let rec = rerun::RecordingStreamBuilder::new("peng").spawn().unwrap();
3291/// let camera = Camera::new((640, 480), 0.1, 100.0, 60.0);
3292/// let use_multi_threading = false;
3293/// log_depth_image(&rec, &camera, use_multi_threading).unwrap();
3294/// let use_multi_threading = true;
3295/// log_depth_image(&rec, &camera, use_multi_threading).unwrap();
3296/// ```
3297pub fn log_depth_image(
3298    rec: &rerun::RecordingStream,
3299    cam: &Camera,
3300    use_multi_threading: bool,
3301) -> Result<(), SimulationError> {
3302    let (width, height) = (cam.resolution.0, cam.resolution.1);
3303    let (min_depth, max_depth) = (cam.near, cam.far);
3304    let depth_image = &cam.depth_buffer;
3305    let mut image: rerun::external::ndarray::Array<u8, _> =
3306        rerun::external::ndarray::Array::zeros((height, width, 3));
3307    let depth_range = max_depth - min_depth;
3308    let scale_factor = 255.0 / depth_range;
3309    if use_multi_threading {
3310        const CHUNK_SIZE: usize = 32;
3311        image
3312            .as_slice_mut()
3313            .expect("Failed to get mutable slice of image")
3314            .par_chunks_exact_mut(CHUNK_SIZE * 3)
3315            .enumerate()
3316            .for_each(|(chunk_index, chunk)| {
3317                let start_index = chunk_index * 32;
3318                for (i, pixel) in chunk.chunks_exact_mut(3).enumerate() {
3319                    let idx = start_index + i;
3320                    let depth = depth_image[idx];
3321                    let color = if depth.is_finite() {
3322                        let normalized_depth =
3323                            ((depth - min_depth) * scale_factor).clamp(0.0, 255.0);
3324                        color_map_fn(normalized_depth)
3325                    } else {
3326                        (0, 0, 0)
3327                    };
3328                    (pixel[0], pixel[1], pixel[2]) = color;
3329                }
3330            });
3331    } else {
3332        for (index, pixel) in image
3333            .as_slice_mut()
3334            .expect("Failed to get mutable slice of image")
3335            .chunks_exact_mut(3)
3336            .enumerate()
3337        {
3338            let depth = depth_image[index];
3339            let color = if depth.is_finite() {
3340                let normalized_depth = ((depth - min_depth) * scale_factor).clamp(0.0, 255.0);
3341                color_map_fn(normalized_depth)
3342            } else {
3343                (0, 0, 0)
3344            };
3345            (pixel[0], pixel[1], pixel[2]) = color;
3346        }
3347    }
3348    let rerun_image = rerun::Image::from_color_model_and_tensor(rerun::ColorModel::RGB, image)
3349        .map_err(|e| SimulationError::OtherError(format!("Failed to create rerun image: {e}")))?;
3350    rec.log("world/quad/cam/depth", &rerun_image)?;
3351    Ok(())
3352}
3353/// creates pinhole camera
3354/// # Arguments
3355/// * `rec` - The rerun::RecordingStream instance
3356/// * `cam` - The camera object
3357/// * `cam_position` - The position vector of the camera (aligns with the quad)
3358/// * `cam_orientation` - The orientation quaternion of quad
3359/// * `cam_transform` - The transform matrix between quad and camera alignment
3360/// # Errors    
3361/// * If the data cannot be logged to the recording stream
3362/// # Example
3363/// ```no_run
3364/// use peng_quad::{log_pinhole_depth, Camera};
3365/// use nalgebra::{Vector3, UnitQuaternion};
3366/// let rec = rerun::RecordingStreamBuilder::new("peng").spawn().unwrap();
3367/// let depth_image = vec![ 0.0f32 ; 640 * 480];
3368/// let cam_position = Vector3::new(0.0,0.0,0.0);
3369/// let cam_orientation = UnitQuaternion::identity();
3370/// let cam_transform = [0.0, 0.0, 1.0, -1.0, 0.0, 0.0, 0.0, -1.0, 0.0];
3371/// let camera = Camera::new((800, 600), 60.0, 0.1, 100.0);
3372/// log_pinhole_depth(&rec, &camera, cam_position, cam_orientation, cam_transform).unwrap();
3373/// ```
3374pub fn log_pinhole_depth(
3375    rec: &rerun::RecordingStream,
3376    cam: &Camera,
3377    cam_position: Vector3<f32>,
3378    cam_orientation: UnitQuaternion<f32>,
3379    cam_transform: [f32; 9],
3380) -> Result<(), SimulationError> {
3381    let depth_image = &cam.depth_buffer;
3382    let (width, height) = cam.resolution;
3383    let pinhole_camera = rerun::Pinhole::from_focal_length_and_resolution(
3384        (cam.horizontal_focal_length, cam.vertical_focal_length),
3385        (width as f32, height as f32),
3386    )
3387    .with_camera_xyz(rerun::components::ViewCoordinates::RDF)
3388    .with_resolution((width as f32, height as f32))
3389    .with_principal_point((width as f32 / 2.0, height as f32 / 2.0));
3390    let rotated_camera_orientation = UnitQuaternion::from_rotation_matrix(
3391        &(cam_orientation.to_rotation_matrix()
3392            * Rotation3::from_matrix_unchecked(Matrix3::from_row_slice(&cam_transform))),
3393    );
3394    let cam_transform = rerun::Transform3D::from_translation_rotation(
3395        rerun::Vec3D::new(cam_position.x, cam_position.y, cam_position.z),
3396        rerun::Quaternion::from_xyzw([
3397            rotated_camera_orientation.i,
3398            rotated_camera_orientation.j,
3399            rotated_camera_orientation.k,
3400            rotated_camera_orientation.w,
3401        ]),
3402    );
3403    rec.log("world/quad/cam", &cam_transform)?;
3404    rec.log("world/quad/cam", &pinhole_camera)?;
3405    let depth_image_rerun =
3406        rerun::external::ndarray::Array::from_shape_vec((height, width), depth_image.to_vec())
3407            .unwrap();
3408    rec.log(
3409        "world/quad/cam/rerun_depth",
3410        &rerun::DepthImage::try_from(depth_image_rerun)
3411            .unwrap()
3412            .with_meter(1.0),
3413    )?;
3414
3415    Ok(())
3416}
3417/// turbo color map function
3418/// # Arguments
3419/// * `gray` - The gray value in the range [0, 255]
3420/// # Returns
3421/// * The RGB color value in the range [0, 255]
3422/// # Example
3423/// ```
3424/// use peng_quad::color_map_fn;
3425/// let color = color_map_fn(128.0);
3426/// ```
3427#[inline]
3428pub fn color_map_fn(gray: f32) -> (u8, u8, u8) {
3429    let x = gray / 255.0;
3430    let r = (34.61
3431        + x * (1172.33 - x * (10793.56 - x * (33300.12 - x * (38394.49 - x * 14825.05)))))
3432        .clamp(0.0, 255.0) as u8;
3433    let g = (23.31 + x * (557.33 + x * (1225.33 - x * (3574.96 - x * (1073.77 + x * 707.56)))))
3434        .clamp(0.0, 255.0) as u8;
3435    let b = (27.2 + x * (3211.1 - x * (15327.97 - x * (27814.0 - x * (22569.18 - x * 6838.66)))))
3436        .clamp(0.0, 255.0) as u8;
3437    (r, g, b)
3438}
3439
3440/// Fast square root function
3441/// # Arguments
3442/// * `x` - The input value
3443/// # Returns
3444/// * The square root of the input value
3445#[inline(always)]
3446fn fast_sqrt(x: f32) -> f32 {
3447    let i = x.to_bits();
3448    let i = 0x1fbd1df5 + (i >> 1);
3449    f32::from_bits(i)
3450}