use std::f64::consts::PI;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SolverType {
    ForwardEuler,
    BackwardEuler,
    Trapezoidal,
    RungeKutta,
}

#[derive(Debug, Clone, Copy, PartialEq)]
pub struct LowPassFilterConfig {
    pub cut_off_frequency_rad_per_sec: f64,
    pub damping_ratio: f64,
    pub initial_condition: f64,
    pub solver_type: SolverType,
}

impl Default for LowPassFilterConfig {
    fn default() -> Self {
        Self {
            cut_off_frequency_rad_per_sec: 20.0,
            damping_ratio: 1.0,
            initial_condition: 0.0,
            solver_type: SolverType::RungeKutta,
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq)]
pub struct LowPassFilterState {
    pub feedback_state: f64,
    pub filtered_output: f64,
}

#[derive(Debug, Clone, PartialEq)]
pub struct SecondOrderLowPassFilter {
    config: LowPassFilterConfig,
    state: LowPassFilterState,
}

impl SecondOrderLowPassFilter {
    pub fn new(config: LowPassFilterConfig) -> Self {
        Self {
            state: LowPassFilterState {
                feedback_state: config.initial_condition,
                filtered_output: config.initial_condition,
            },
            config,
        }
    }

    pub const fn state(&self) -> LowPassFilterState {
        self.state
    }

    pub fn step(&mut self, input: f64, dt: f64) -> (f64, f64) {
        let dt = dt.max(0.0);
        let wn = self.config.cut_off_frequency_rad_per_sec;
        let zeta = self.config.damping_ratio;
        let deriv = |feedback: f64, output: f64| {
            let input_error = input - output;
            let corrected_error = input_error - (2.0 * zeta * feedback);
            let feedback_dot = wn * corrected_error;
            let output_dot = wn * feedback;
            (feedback_dot, output_dot)
        };

        let (feedback_out, output_out, feedback_next, output_next) = match self.config.solver_type {
            SolverType::ForwardEuler => {
                let (dq, dy) = deriv(self.state.feedback_state, self.state.filtered_output);
                (
                    self.state.feedback_state,
                    self.state.filtered_output,
                    self.state.feedback_state + dt * dq,
                    self.state.filtered_output + dt * dy,
                )
            }
            SolverType::BackwardEuler => {
                let (dq, dy) = deriv(self.state.feedback_state, self.state.filtered_output);
                let feedback = self.state.feedback_state + dt * dq;
                let output = self.state.filtered_output + dt * dy;
                (feedback, output, feedback, output)
            }
            SolverType::Trapezoidal => {
                let (dq, dy) = deriv(self.state.feedback_state, self.state.filtered_output);
                let feedback = self.state.feedback_state + (dt * 0.5) * dq;
                let output = self.state.filtered_output + (dt * 0.5) * dy;
                (
                    feedback,
                    output,
                    feedback + (dt * 0.5) * dq,
                    output + (dt * 0.5) * dy,
                )
            }
            SolverType::RungeKutta => {
                let (k1_q, k1_y) = deriv(self.state.feedback_state, self.state.filtered_output);
                let (k2_q, k2_y) = deriv(
                    self.state.feedback_state + 0.5 * dt * k1_q,
                    self.state.filtered_output + 0.5 * dt * k1_y,
                );
                let (k3_q, k3_y) = deriv(
                    self.state.feedback_state + 0.5 * dt * k2_q,
                    self.state.filtered_output + 0.5 * dt * k2_y,
                );
                let (k4_q, k4_y) = deriv(
                    self.state.feedback_state + dt * k3_q,
                    self.state.filtered_output + dt * k3_y,
                );
                let feedback = self.state.feedback_state
                    + (dt / 6.0) * (k1_q + 2.0 * k2_q + 2.0 * k3_q + k4_q);
                let output = self.state.filtered_output
                    + (dt / 6.0) * (k1_y + 2.0 * k2_y + 2.0 * k3_y + k4_y);
                (feedback, output, feedback, output)
            }
        };

        self.state.feedback_state = feedback_next;
        self.state.filtered_output = output_next;

        let output_derivative = wn * feedback_out;
        (output_out, output_derivative)
    }

    pub fn run(&mut self, inputs: &[f64], timestamps: Option<&[f64]>) -> (Vec<f64>, Vec<f64>) {
        self.reset();
        let mut outputs = Vec::with_capacity(inputs.len());
        let mut derivatives = Vec::with_capacity(inputs.len());

        for (idx, input) in inputs.iter().copied().enumerate() {
            let dt = match timestamps {
                Some(values) if idx > 0 => values[idx] - values[idx - 1],
                _ => 0.01,
            };
            let (output, derivative) = self.step(input, dt);
            outputs.push(output);
            derivatives.push(derivative);
        }

        (outputs, derivatives)
    }

    pub fn reset(&mut self) {
        self.state.feedback_state = self.config.initial_condition;
        self.state.filtered_output = self.config.initial_condition;
    }
}

impl Default for SecondOrderLowPassFilter {
    fn default() -> Self {
        Self::new(LowPassFilterConfig::default())
    }
}

#[derive(Debug, Clone, Copy, PartialEq)]
pub struct PIDConfig {
    pub proportional_gain: f64,
    pub integral_gain: f64,
    pub derivative_gain: f64,
    pub output_limits: Option<(f64, f64)>,
}

impl Default for PIDConfig {
    fn default() -> Self {
        Self {
            proportional_gain: 2.0,
            integral_gain: 0.0,
            derivative_gain: 0.02,
            output_limits: Some((-7.0, 7.0)),
        }
    }
}

#[derive(Debug, Clone, PartialEq)]
pub struct PIDController {
    pub config: PIDConfig,
    pub low_pass_filter: SecondOrderLowPassFilter,
    integral: f64,
    previous_velocity: f64,
    previous_timestamp: Option<f64>,
}

impl PIDController {
    pub fn new(config: PIDConfig, filter_config: LowPassFilterConfig) -> Self {
        Self {
            config,
            low_pass_filter: SecondOrderLowPassFilter::new(filter_config),
            integral: 0.0,
            previous_velocity: 0.0,
            previous_timestamp: None,
        }
    }

    pub fn compute_output(
        &mut self,
        timestamp: f64,
        target_position: f64,
        motor_position: f64,
    ) -> f64 {
        let Some(previous_timestamp) = self.previous_timestamp else {
            self.previous_timestamp = Some(timestamp);
            return 0.0;
        };

        let dt = (timestamp - previous_timestamp).max(0.0);
        let max_error = PI / 3.0;
        let error = (target_position - motor_position).clamp(-max_error, max_error);

        let proportional = self.config.proportional_gain * error;
        self.integral += error * dt;
        let integral = self.config.integral_gain * self.integral;
        let (low, high) = self.config.output_limits.unwrap_or((-100.0, 100.0));
        let integral = integral.clamp(low, high);

        let (damping_feedback, _) = self.low_pass_filter.step(self.previous_velocity, dt);
        let derivative = self.config.derivative_gain * damping_feedback;
        let output = (proportional + integral - derivative).clamp(low, high);

        self.previous_velocity = output;
        self.previous_timestamp = Some(timestamp);
        output
    }

    pub const fn integral(&self) -> f64 {
        self.integral
    }

    pub const fn previous_velocity(&self) -> f64 {
        self.previous_velocity
    }

    pub fn reset(&mut self) {
        self.integral = 0.0;
        self.previous_velocity = 0.0;
        self.previous_timestamp = None;
        self.low_pass_filter.reset();
    }
}

impl Default for PIDController {
    fn default() -> Self {
        Self::new(PIDConfig::default(), LowPassFilterConfig::default())
    }
}
