#[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)
    }
}

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

    #[test]
    fn test_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));
    }
}