#[derive(Debug, Clone, Copy, PartialEq)]
pub enum TryStep {
    Accepted(f64, f64),
    NotYetAccepted(f64),
}

impl TryStep {
    pub fn extract(&self) -> f64 {
        match self {
            TryStep::Accepted(h, _) => *h,
            TryStep::NotYetAccepted(h) => *h,
        }
    }

    pub fn is_accepted(&self) -> bool {
        matches!(self, TryStep::Accepted(_, _))
    }

    pub fn reset(&mut self) -> Result<TryStep, &str> {
        match self {
            TryStep::Accepted(_, h) => Ok(TryStep::NotYetAccepted(*h)),
            TryStep::NotYetAccepted(_) => Err("Cannot reset a NotYetAccepted TryStep"),
        }
    }
}

pub trait Controller<const D: usize> {
    fn determine_step(&mut self, h: f64, err: f64) -> TryStep;
}

#[derive(Debug, Clone, Copy)]
pub struct PIController {
    pub alpha: f64,
    pub beta: f64,
    pub factor_c1: f64,
    pub factor_c2: f64,
    pub factor_old: f64,
    pub h_max: f64,
    pub safety_factor: f64,
    pub next_step_guess: TryStep,
}

impl<const D: usize> Controller<D> for PIController {
    /// Determines if the previously run step size and error were valid or not. Either way, it also
    /// returns what the next step size should be
    fn determine_step(&mut self, prev_step: f64, err: f64) -> TryStep {
        let factor_11 = err.powf(self.alpha);
        let factor = self.factor_c2.max(
            self.factor_c1
                .min(factor_11 * self.factor_old.powf(-self.beta) / self.safety_factor),
        );
        if err <= 1.0 {
            let mut h = prev_step / factor;
            // Accept the stepsize and provide what the next step size should be
            self.factor_old = err.max(1.0e-4);
            if h.abs() > self.h_max {
                // If the step goes past the maximum allowed, though, we shrink it
                h = self.h_max.copysign(h);
            }
            TryStep::Accepted(prev_step, h)
        } else {
            // Reject the stepsize and propose a smaller one for the current step
            TryStep::NotYetAccepted(prev_step / (self.factor_c1.min(factor_11 / self.safety_factor)))
        }
    }
}

impl PIController {
    pub fn new(
        alpha: f64,
        beta: f64,
        max_factor: f64,
        min_factor: f64,
        h_max: f64,
        safety_factor: f64,
        initial_h: f64,
    ) -> Self {
        Self {
            alpha,
            beta,
            factor_c1: 1.0 / min_factor,
            factor_c2: 1.0 / max_factor,
            factor_old: 1.0e-4,
            h_max: h_max.abs(),
            safety_factor,
            next_step_guess: TryStep::NotYetAccepted(initial_h),
        }
    }
}

impl Default for PIController {
    fn default() -> Self {
        Self::new(0.17, 0.04, 10.0, 0.2, 100000.0, 0.9, 1e-4)
    }
}

/// PID (Proportional-Integral-Derivative) step size controller
///
/// The PID controller is an advanced adaptive time-stepping controller that provides
/// better stability than the PI controller, especially for difficult or oscillatory problems.
///
/// # Mathematical Formulation
///
/// The PID controller determines the next step size based on error estimates from the
/// current and previous two steps:
///
/// ```text
/// h_{n+1} = h_n * safety * (ε_n)^(-β₁) * (ε_{n-1})^(-β₂) * (h_n/h_{n-1})^(-β₃)
/// ```
///
/// Where:
/// - ε_n = normalized error estimate at current step
/// - ε_{n-1} = normalized error estimate at previous step
/// - β₁ = proportional coefficient (controls reaction to current error)
/// - β₂ = integral coefficient (controls reaction to error history)
/// - β₃ = derivative coefficient (controls reaction to error rate of change)
///
/// # When to Use
///
/// Prefer PID over PI when:
/// - Problem exhibits step size oscillation with PI controller
/// - Working with stiff or near-stiff problems
/// - Need smoother step size evolution
/// - Standard in production solvers (MATLAB, Sundials)
///
/// # Example
///
/// ```ignore
/// use ordinary_diffeq::prelude::*;
///
/// // Default PID controller (conservative coefficients)
/// let controller = PIDController::default();
///
/// // Custom PID controller
/// let controller = PIDController::new(0.49, 0.34, 0.10, 10.0, 0.2, 100.0, 0.9, 1e-4);
///
/// // PID tuned for specific method order (Gustafsson coefficients)
/// let controller = PIDController::for_order(5);
/// ```
#[derive(Debug, Clone, Copy)]
pub struct PIDController {
    // PID Coefficients
    pub beta1: f64,  // Proportional: reaction to current error
    pub beta2: f64,  // Integral: reaction to error history
    pub beta3: f64,  // Derivative: reaction to error rate of change

    // Constraints
    pub factor_c1: f64,      // 1/min_factor (maximum step decrease)
    pub factor_c2: f64,      // 1/max_factor (maximum step increase)
    pub h_max: f64,          // Maximum allowed step size
    pub safety_factor: f64,  // Safety factor (typically 0.9)

    // Error history for PID control
    pub err_old: f64,     // ε_{n-1}: previous step error
    pub err_older: f64,   // ε_{n-2}: error two steps ago
    pub h_old: f64,       // h_{n-1}: previous step size

    // Next step guess
    pub next_step_guess: TryStep,
}

impl<const D: usize> Controller<D> for PIDController {
    /// Determines if the previously run step was acceptable and computes the next step size
    /// using PID control theory
    fn determine_step(&mut self, h: f64, err: f64) -> TryStep {
        // Compute PID control factor
        // For first step or when history isn't available, fall back to simpler control
        let factor = if self.err_old <= 0.0 {
            // First step: use only proportional control
            let factor_11 = err.powf(self.beta1);
            self.factor_c2.max(
                self.factor_c1.min(factor_11 / self.safety_factor)
            )
        } else if self.err_older <= 0.0 {
            // Second step: use PI control (proportional + integral)
            let factor_11 = err.powf(self.beta1);
            let factor_12 = self.err_old.powf(-self.beta2);
            self.factor_c2.max(
                self.factor_c1.min(factor_11 * factor_12 / self.safety_factor)
            )
        } else {
            // Full PID control (proportional + integral + derivative)
            let factor_11 = err.powf(self.beta1);
            let factor_12 = self.err_old.powf(-self.beta2);
            // Derivative term uses ratio of consecutive step sizes
            let factor_13 = if self.h_old > 0.0 {
                (h / self.h_old).powf(-self.beta3)
            } else {
                1.0
            };
            self.factor_c2.max(
                self.factor_c1.min(factor_11 * factor_12 * factor_13 / self.safety_factor)
            )
        };

        if err <= 1.0 {
            // Step accepted
            let mut h_next = h / factor;

            // Update error history for next step
            self.err_older = self.err_old;
            self.err_old = err.max(1.0e-4);  // Prevent very small values
            self.h_old = h;

            // Apply maximum step size limit
            if h_next.abs() > self.h_max {
                h_next = self.h_max.copysign(h_next);
            }

            TryStep::Accepted(h, h_next)
        } else {
            // Step rejected - propose smaller step
            // Use only proportional control for rejection (more aggressive)
            let factor_11 = err.powf(self.beta1);
            let h_next = h / (self.factor_c1.min(factor_11 / self.safety_factor));

            // Note: Don't update history on rejection
            TryStep::NotYetAccepted(h_next)
        }
    }
}

impl PIDController {
    /// Create a new PID controller with custom parameters
    ///
    /// # Arguments
    ///
    /// * `beta1` - Proportional coefficient (typically 0.3-0.5)
    /// * `beta2` - Integral coefficient (typically 0.04-0.1)
    /// * `beta3` - Derivative coefficient (typically 0.01-0.05)
    /// * `max_factor` - Maximum step size increase factor (typically 10.0)
    /// * `min_factor` - Maximum step size decrease factor (typically 0.2)
    /// * `h_max` - Maximum allowed step size
    /// * `safety_factor` - Safety factor (typically 0.9)
    /// * `initial_h` - Initial step size guess
    pub fn new(
        beta1: f64,
        beta2: f64,
        beta3: f64,
        max_factor: f64,
        min_factor: f64,
        h_max: f64,
        safety_factor: f64,
        initial_h: f64,
    ) -> Self {
        Self {
            beta1,
            beta2,
            beta3,
            factor_c1: 1.0 / min_factor,
            factor_c2: 1.0 / max_factor,
            h_max: h_max.abs(),
            safety_factor,
            err_old: 0.0,      // No history initially
            err_older: 0.0,    // No history initially
            h_old: 0.0,        // No history initially
            next_step_guess: TryStep::NotYetAccepted(initial_h),
        }
    }

    /// Create a PID controller with coefficients scaled for a specific method order
    ///
    /// Uses the Gustafsson coefficients scaled by order:
    /// - β₁ = 0.49 / (order + 1)
    /// - β₂ = 0.34 / (order + 1)
    /// - β₃ = 0.10 / (order + 1)
    ///
    /// # Arguments
    ///
    /// * `order` - Order of the integration method (e.g., 5 for DP5, 7 for Vern7)
    pub fn for_order(order: usize) -> Self {
        let k_plus_1 = (order + 1) as f64;
        Self::new(
            0.49 / k_plus_1,  // beta1: proportional
            0.34 / k_plus_1,  // beta2: integral
            0.10 / k_plus_1,  // beta3: derivative
            10.0,              // max_factor
            0.2,               // min_factor
            100000.0,          // h_max
            0.9,               // safety_factor
            1e-4,              // initial_h
        )
    }

    /// Reset the controller's error history
    ///
    /// Useful when switching algorithms or restarting integration
    pub fn reset(&mut self) {
        self.err_old = 0.0;
        self.err_older = 0.0;
        self.h_old = 0.0;
    }
}

impl Default for PIDController {
    /// Default PID controller using conservative coefficients
    ///
    /// Uses conservative PID coefficients that provide stable performance
    /// across a wide range of problems:
    /// - β₁ = 0.07 (proportional)
    /// - β₂ = 0.04 (integral)
    /// - β₃ = 0.01 (derivative)
    ///
    /// These values are appropriate for typical ODE methods of order 5-7.
    /// For method-specific tuning, use `PIDController::for_order(order)` instead.
    fn default() -> Self {
        Self::new(
            0.07,      // beta1 (proportional)
            0.04,      // beta2 (integral)
            0.01,      // beta3 (derivative)
            10.0,      // max_factor
            0.2,       // min_factor
            100000.0,  // h_max
            0.9,       // safety_factor
            1e-4,      // initial_h
        )
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_pi_controller_creation() {
        let controller = PIController::new(0.17, 0.04, 10.0, 0.2, 10.0, 0.9, 1e-4);

        assert!(controller.alpha == 0.17);
        assert!(controller.beta == 0.04);
        assert!(controller.factor_c1 == 1.0 / 0.2);
        assert!(controller.factor_c2 == 1.0 / 10.0);
        assert!(controller.factor_old == 1.0e-4);
        assert!(controller.h_max == 10.0);
        assert!(controller.safety_factor == 0.9);
        assert!(controller.next_step_guess == TryStep::NotYetAccepted(1e-4));
    }

    #[test]
    fn test_pid_controller_creation() {
        let controller = PIDController::new(0.49, 0.34, 0.10, 10.0, 0.2, 10.0, 0.9, 1e-4);

        assert_eq!(controller.beta1, 0.49);
        assert_eq!(controller.beta2, 0.34);
        assert_eq!(controller.beta3, 0.10);
        assert_eq!(controller.factor_c1, 1.0 / 0.2);
        assert_eq!(controller.factor_c2, 1.0 / 10.0);
        assert_eq!(controller.h_max, 10.0);
        assert_eq!(controller.safety_factor, 0.9);
        assert_eq!(controller.err_old, 0.0);
        assert_eq!(controller.err_older, 0.0);
        assert_eq!(controller.h_old, 0.0);
    }

    #[test]
    fn test_pid_for_order() {
        let controller = PIDController::for_order(5);

        // For order 5, k+1 = 6
        assert!((controller.beta1 - 0.49 / 6.0).abs() < 1e-10);
        assert!((controller.beta2 - 0.34 / 6.0).abs() < 1e-10);
        assert!((controller.beta3 - 0.10 / 6.0).abs() < 1e-10);
    }

    #[test]
    fn test_pid_default() {
        let controller = PIDController::default();

        // Default uses conservative PID coefficients
        assert_eq!(controller.beta1, 0.07);
        assert_eq!(controller.beta2, 0.04);
        assert_eq!(controller.beta3, 0.01);  // True PID with derivative term
    }

    #[test]
    fn test_pid_reset() {
        let mut controller = PIDController::default();

        // Simulate some history
        controller.err_old = 0.5;
        controller.err_older = 0.3;
        controller.h_old = 0.01;

        controller.reset();

        assert_eq!(controller.err_old, 0.0);
        assert_eq!(controller.err_older, 0.0);
        assert_eq!(controller.h_old, 0.0);
    }

    #[test]
    fn test_pi_vs_pid_smooth_problem() {
        // For smooth problems, PI and PID should perform similarly
        // Test with exponential decay: y' = -y

        let mut pi = PIController::default();
        let mut pid = PIDController::default();

        // Simulate a sequence of small errors (smooth problem)
        let errors = vec![0.8, 0.6, 0.5, 0.45, 0.4, 0.35, 0.3];
        let h = 0.01;

        let mut pi_steps = Vec::new();
        let mut pid_steps = Vec::new();

        for &err in &errors {
            let mut pi_result = <PIController as Controller<1>>::determine_step(&mut pi, h, err);
            let mut pid_result = <PIDController as Controller<1>>::determine_step(&mut pid, h, err);

            if pi_result.is_accepted() {
                pi_steps.push(pi_result.extract());
                pi.next_step_guess = pi_result.reset().unwrap();
            }

            if pid_result.is_accepted() {
                pid_steps.push(pid_result.extract());
                pid.next_step_guess = pid_result.reset().unwrap();
            }
        }

        // Both should accept all steps for this smooth sequence
        assert_eq!(pi_steps.len(), errors.len());
        assert_eq!(pid_steps.len(), errors.len());

        // Step sizes should be reasonably similar (within 20%)
        // PID may differ slightly due to derivative term
        for (pi_h, pid_h) in pi_steps.iter().zip(pid_steps.iter()) {
            let relative_diff = ((pi_h - pid_h) / pi_h).abs();
            assert!(
                relative_diff < 0.2,
                "Step sizes differ by more than 20%: PI={}, PID={}",
                pi_h,
                pid_h
            );
        }
    }

    #[test]
    fn test_pid_bootstrap() {
        // Test that PID progressively uses P → PI → PID as history builds
        let mut pid = PIDController::new(0.49, 0.34, 0.10, 10.0, 0.2, 100.0, 0.9, 0.01);

        let h = 0.01;
        let err1 = 0.5;
        let err2 = 0.4;
        let err3 = 0.3;

        // First step: should use only proportional (beta1)
        assert_eq!(pid.err_old, 0.0);
        assert_eq!(pid.err_older, 0.0);
        let step1 = <PIDController as Controller<1>>::determine_step(&mut pid, h, err1);
        assert!(step1.is_accepted());

        // After first step, err_old is updated but err_older is still 0
        assert!(pid.err_old > 0.0);
        assert_eq!(pid.err_older, 0.0);

        // Second step: should use PI (beta1 and beta2)
        let step2 = <PIDController as Controller<1>>::determine_step(&mut pid, h, err2);
        assert!(step2.is_accepted());

        // After second step, both err_old and err_older should be set
        assert!(pid.err_old > 0.0);
        assert!(pid.err_older > 0.0);
        assert!(pid.h_old > 0.0);

        // Third step: should use full PID (beta1, beta2, and beta3)
        let step3 = <PIDController as Controller<1>>::determine_step(&mut pid, h, err3);
        assert!(step3.is_accepted());
    }

    #[test]
    fn test_pid_step_rejection() {
        // Test that error history is NOT updated after rejection
        let mut pid = PIDController::default();

        let h = 0.01;

        // First, accept a step to build history
        let err_good = 0.5;
        let step1 = <PIDController as Controller<1>>::determine_step(&mut pid, h, err_good);
        assert!(step1.is_accepted());

        let err_old_before = pid.err_old;
        let err_older_before = pid.err_older;
        let h_old_before = pid.h_old;

        // Now reject a step with large error
        let err_bad = 2.0;
        let step2 = <PIDController as Controller<1>>::determine_step(&mut pid, h, err_bad);
        assert!(!step2.is_accepted());

        // History should NOT have changed after rejection
        assert_eq!(pid.err_old, err_old_before);
        assert_eq!(pid.err_older, err_older_before);
        assert_eq!(pid.h_old, h_old_before);
    }

    #[test]
    fn test_pid_vs_pi_oscillatory() {
        // Test on oscillatory error pattern (simulating difficult problem)
        // True PID (with derivative term) should provide smoother step size evolution

        let mut pi = PIController::default();
        // Use actual PID with non-zero beta3 (Gustafsson coefficients for order 5)
        let mut pid = PIDController::for_order(5);

        // Simulate oscillatory error pattern
        let errors = vec![0.8, 0.3, 0.9, 0.2, 0.85, 0.25, 0.8, 0.3];
        let h = 0.01;

        let mut pi_ratios = Vec::new();
        let mut pid_ratios = Vec::new();

        let mut pi_h_prev = h;
        let mut pid_h_prev = h;

        for &err in &errors {
            let mut pi_result = <PIController as Controller<1>>::determine_step(&mut pi, h, err);
            let mut pid_result = <PIDController as Controller<1>>::determine_step(&mut pid, h, err);

            if pi_result.is_accepted() {
                let pi_h_next = pi_result.reset().unwrap().extract();
                pi_ratios.push((pi_h_next / pi_h_prev).ln().abs());
                pi_h_prev = pi_h_next;
                pi.next_step_guess = TryStep::NotYetAccepted(pi_h_next);
            }

            if pid_result.is_accepted() {
                let pid_h_next = pid_result.reset().unwrap().extract();
                pid_ratios.push((pid_h_next / pid_h_prev).ln().abs());
                pid_h_prev = pid_h_next;
                pid.next_step_guess = TryStep::NotYetAccepted(pid_h_next);
            }
        }

        // Compute standard deviation of log step size ratios
        let pi_mean: f64 = pi_ratios.iter().sum::<f64>() / pi_ratios.len() as f64;
        let pi_variance: f64 = pi_ratios
            .iter()
            .map(|r| (r - pi_mean).powi(2))
            .sum::<f64>()
            / pi_ratios.len() as f64;
        let pi_std = pi_variance.sqrt();

        let pid_mean: f64 = pid_ratios.iter().sum::<f64>() / pid_ratios.len() as f64;
        let pid_variance: f64 = pid_ratios
            .iter()
            .map(|r| (r - pid_mean).powi(2))
            .sum::<f64>()
            / pid_ratios.len() as f64;
        let pid_std = pid_variance.sqrt();

        // With true PID (non-zero beta3), we expect similar or better stability
        // Allow some tolerance since this is a simple synthetic test
        assert!(
            pid_std <= pi_std * 1.1,
            "PID should not be significantly worse than PI: PI_std={:.3}, PID_std={:.3}",
            pi_std,
            pid_std
        );
    }
}