use crate::config::PidConfig;
use crate::drone::stacked::DroneState;
use nalgebra as na;

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

pub struct ModuleRuntime {
    pub module: ControllerModule,
    pub target_dt: f32, // e.g., 0.01 for 100Hz
    pub accumulated_time: f32,
    pub last_output: na::Vector3<f32>, // Store the last result to pass down the chain
}

impl ControllerModule {
    pub fn process(
        &mut self,
        setpoint: na::Vector3<f32>,
        state: &DroneState,
        dt: f32,
        is_first_layer: bool,
    ) -> na::Vector3<f32> {
        match self {
            ControllerModule::Angle {
                processor,
                max_angle,
            } => {
                // Setpoint is -1.0..1.0, scale it to target Radians
                let target_angles = if is_first_layer {
                    setpoint * *max_angle
                } else {
                    setpoint
                };

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

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

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

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

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

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

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

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