protoflight 0.1.2

use pidsk_controller::{PidControllerf32, PidGainsf32};
use vqm::Quaternionf32;

/// Altitude hold uses a standard **Dual-Ring Cascaded PID Loop**.
///
/// Because a multirotor controls acceleration (via motor thrust) to change altitude,
/// a single PID loop often suffers from severe oscillation or sluggish response.
///
/// So, instead of mapping altitude error directly to throttle, there are two nested loops:
///
/// 1. Outer Loop (Altitude): Takes altitude error and outputs a target vertical speed (climb rate).
/// 2. Inner Loop (Vertical Speed): Takes vertical speed error and outputs the throttle offset.
///
/// Because vertical speed is the derivative of altitude, the outer loop does not need a Dterm,
/// since the inner loop Pterm is effectively the outer loop Dterm.
///
/// The inner loop incorporates some open-loop control via its Kterm (kick), so there is
/// an immediate response when its setpoint changes (no need to wait for the error to accumulate).
///
/// ```text
/// [Target Altitude]
///       │
///       ▼
///   ( + / - ) <─── [Current Altitude]
///       │
///       ▼
/// ┌───────────┐
/// │ P-Loop    │ (Outer Loop)
/// └─────┬─────┘
///       │
/// [Target Vertical Speed]
///       │
///       ▼
///   ( + / - ) <─── [Current Vertical Speed]
///       │
///       ▼
/// ┌───────────┐
/// │ PIDK-Loop │ (Inner Loop)
/// └─────┬─────┘
///       │
///       ▼
/// [Throttle Offset] ──> + Base Throttle ──> Motors
///
/// Since vertical speed is the first derivative of altitude the outer loop can be a pure P-controller.
///
/// 1. No Dterm is required because he inner speed loop inherently acts as the "D" (derivative) term for the outer position loop.
///    (Adding a D-term to the outer loop would mean calculating the derivative of velocity (acceleration), which would introduce severe sensor noise).
/// 2. No Iterm is required because steady-state errors (like gravity dragging the aircraft down) are handled by the inner speed loop's Iterm (adjusting throttle bias).
/// ```
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct AltitudeDualRingPid {
    /// Outer Loop: Inputs Target Altitude -> Output Target Vertical Speed.
    altitude_pid: PidControllerf32,
    /// Inner Loop: Input Target Vertical Speed -> Output Throttle Adjustment.
    speed_pid: PidControllerf32,

    // Core parameters
    /// max vertical speed (climb rate), m/s.
    max_vertical_speed_mps: f32,
    /// maximum allowed thrust adjustment.
    max_throttle_adjustment: f32,
    /// estimated throttle needed to hover.
    hover_throttle: f32,
}

impl Default for AltitudeDualRingPid {
    fn default() -> Self {
        Self::new(0.0)
    }
}

impl AltitudeDualRingPid {
    pub const fn new(hover_throttle: f32) -> Self {
        Self {
            // Initialize height controller (Outer Loop)
            // Only needs Proportional (kp) to map distance error to speed:
            // because the inner loop handles the physics of acceleration,
            // the outer loop only needs Kp to calculate the vertical speed setpoint
            altitude_pid: PidControllerf32::new(1.0),
            // Initialize velocity controller (Inner Loop)
            // Highly reactive: utilizes kp, ki, and kd.
            // TODO: check default PID gains.
            speed_pid: PidControllerf32::with_gains(PidGainsf32 { kp: 2.5, ki: 0.05, kd: 0.05, ks: 0.0, kk: 0.0 }),
            max_vertical_speed_mps: 10.0, // = 36.0 km/h, effectively unlimited
            max_throttle_adjustment: 1.0, // effectively unlimited
            hover_throttle,
        }
    }
}

#[allow(unused)]
impl AltitudeDualRingPid {
    pub fn set_altitude_setpoint(&mut self, altitude_setpoint: f32) {
        self.altitude_pid.set_setpoint(altitude_setpoint);
    }

    pub fn update(&mut self, altitude: f32, vertical_speed: f32, orientation: Quaternionf32, delta_t: f32) -> f32 {
        self.hover_throttle + self.calculate_throttle_offset(altitude, vertical_speed, orientation, delta_t)
    }

    pub fn calculate_throttle_offset(
        &mut self,
        altitude: f32,
        vertical_speed: f32,
        orientation: Quaternionf32,
        delta_t: f32,
    ) -> f32 {
        let cos_tilt = orientation.cos_tilt();
        if cos_tilt < 0.0 {
            // craft is upside down, so cannot make adjustment
            return 0.0;
        }
        // --- STEP 1: Altitude Loop ---
        let vertical_speed_setpoint = self
            .altitude_pid
            .update_sp(altitude) // just call update_sp, since ki and kd are zero.
            .clamp(-self.max_vertical_speed_mps, self.max_vertical_speed_mps);

        // --- STEP 2: Vertical Speed Loop ---
        self.speed_pid.set_setpoint(vertical_speed_setpoint);

        // calculate throttle offset, adjusting for tilt angle.
        let cos_tilt_reciprocal = 1.0 / cos_tilt.clamp(0.1, 1.0);
        let throttle_offset =
            self.speed_pid.update(vertical_speed, delta_t) + self.hover_throttle * (cos_tilt_reciprocal - 1.0);

        throttle_offset.clamp(-self.max_throttle_adjustment, self.max_throttle_adjustment)
    }
    pub fn reset(&mut self) {
        self.altitude_pid.reset();
        self.speed_pid.reset();
    }
}

#[cfg(test)]
mod tests {
    #![allow(clippy::float_cmp)]

    use crate::autopilot::MockMultirotorZ;

    use super::*;

    fn _is_normal<T: Sized + Send + Sync + Unpin>() {}
    fn is_full<T: Sized + Send + Sync + Unpin + Copy + Clone + Default + PartialEq>() {}

    #[test]
    fn normal_types() {
        is_full::<AltitudeDualRingPid>();
    }
    #[test]
    fn test_new() {
        let _altitude_hold = AltitudeDualRingPid::new(0.0);
    }
    #[test]
    fn test_altitude_hold_convergence() {
        let hover_throttle = 0.5; // Steady state mid-point
        let mut controller = AltitudeDualRingPid::new(hover_throttle);
        let mut multirotor = MockMultirotorZ::new(hover_throttle);

        // --- BALANCED TUNING FOR UNIT SIMULATION ---
        controller.altitude_pid.set_gains(PidGainsf32 { kp: 0.29, ki: 0.0, kd: 0.0, ks: 0.0, kk: 0.0 });
        // Strong P braking, minor I, strong D damping, no kick.
        controller.speed_pid.set_gains(PidGainsf32 { kp: 1.0, ki: 0.05, kd: 0.1, ks: 0.0, kk: 0.0 });

        let altitude_setpoint = 5.0; // Want to climb to 5 meters
        controller.set_altitude_setpoint(altitude_setpoint);
        let delta_t = 0.001; // 1kHz execution speed
        let orientation = Quaternionf32::default();

        // Simulate 15 seconds of real-time flight execution (15,000 loop cycles)
        let mut converged = false;
        for _ in 0..15_000 {
            let throttle = controller.update(multirotor.altitude, multirotor.vertical_speed, orientation, delta_t);

            multirotor.step(throttle, delta_t);

            // Check if we reached the target closely and stabilized vertical velocity
            if (multirotor.altitude - altitude_setpoint).abs() < 0.05 && multirotor.vertical_speed.abs() < 0.01 {
                converged = true;
                break;
            }
        }
        assert!(
            converged,
            "Multirotor failed to settle at target altitude. Final Alt: {}, Vel: {}",
            multirotor.altitude, multirotor.vertical_speed
        );
    }
}