use nalgebra::SVector;
use roots::{find_root_brent, SimpleConvergency};

use super::callback::Callback;
use super::controller::{Controller, PIController};
use super::integrator::Integrator;
use super::ode::ODE;

#[derive(Clone)]
pub struct Problem<'a, const D: usize, S, P>
where
    S: Integrator<D>,
{
    ode: ODE<'a, D, P>,
    integrator: S,
    controller: PIController,
    callbacks: Vec<Callback<'a, D, P>>,
}

impl<'a, const D: usize, S, P> Problem<'a, D, S, P>
where
    S: Integrator<D> + Copy,
{
    pub fn new(ode: ODE<'a, D, P>, integrator: S, controller: PIController) -> Self {
        Problem {
            ode: ode,
            integrator: integrator,
            controller: controller,
            callbacks: Vec::new(),
        }
    }
    pub fn solve(&mut self) -> Solution<S, D> {
        let mut convergency = SimpleConvergency { eps: 1e-10, max_iter: 1000 };
        let mut times: Vec<f64> = vec![self.ode.t];
        let mut states: Vec<SVector<f64, D>> = vec![self.ode.y];
        let mut dense_coefficients: Vec<Vec<SVector<f64, D>>> = Vec::new();
        let mut step: f64 = self.controller.old_h;
        let (mut new_y, mut err_option, _) = self.integrator.step(&self.ode, 0.0);
        while self.ode.t < self.ode.t_end {
            let mut dense_option: Option<Vec<SVector<f64, D>>> = None;
            if S::ADAPTIVE {
                let mut err = err_option.unwrap();
                let mut accepted: bool = false;
                while !accepted {
                    // Try a step and if that isn't acceptable, then change the step until it is
                    (accepted, step) = <PIController as Controller<D>>::determine_step(
                        &mut self.controller,
                        step,
                        err,
                    );
                    (new_y, err_option, dense_option) = self.integrator.step(&self.ode, step);
                    err = err_option.unwrap();
                }
                self.controller.old_h = step;
                self.controller.h_max = self
                    .controller
                    .h_max
                    .min(self.ode.t_end - self.ode.t - step);
            } else {
                // If fixed time step just step forward one step
                (new_y, _, dense_option) = self.integrator.step(&self.ode, step);
            }
            if self.callbacks.len() > 0 {
                // Check for events occurring
                for callback in &self.callbacks {
                    if (callback.event)(self.ode.t, self.ode.y, &self.ode.params)
                        * (callback.event)(self.ode.t + step, new_y, &self.ode.params)
                        < 0.0
                    {
                        // If the event crossed zero, then find the root
                        let f = |test_t| {
                            let test_y = self.integrator.step(&self.ode, test_t).0;
                            (callback.event)(self.ode.t + test_t, test_y, &self.ode.params)
                        };
                        let root = find_root_brent(0.0, step, &f, &mut convergency).unwrap();
                        step = root;
                        (new_y, _, dense_option) = self.integrator.step(&self.ode, step);
                        (callback.effect)(&mut self.ode);
                    }
                }
            }
            self.ode.y = new_y;
            self.ode.t += step;
            times.push(self.ode.t);
            states.push(self.ode.y);
            // TODO: Implement third order interpolation for non-dense algorithms
            dense_coefficients.push(dense_option.unwrap());
        }
        Solution {
            integrator: self.integrator,
            times,
            states,
            dense: dense_coefficients,
        }
    }

    pub fn with_callback(mut self, callback: Callback<'a, D, P>) -> Self {
        self.callbacks.push(callback);
        Self {
            ode: self.ode,
            integrator: self.integrator,
            controller: self.controller,
            callbacks: self.callbacks,
        }
    }
}

pub struct Solution<S, const D: usize>
where
    S: Integrator<D>,
{
    pub integrator: S,
    pub times: Vec<f64>,
    pub states: Vec<SVector<f64, D>>,
    pub dense: Vec<Vec<SVector<f64, D>>>,
}

impl<S, const D: usize> Solution<S, D>
where
    S: Integrator<D>,
{
    pub fn interpolate(&self, t: f64) -> SVector<f64, D> {
        // First check that the t is within bounds
        let last = self.times.last().unwrap();
        let first = self.times.first().unwrap();

        // TODO: Improve these errors
        let mut times = self.times.clone();
        if *first > *last {
            times.reverse();
        }
        if t < *first || t > *last {
            panic!();
        }

        // Then find the two t values closest to the desired t
        let mut end_index: usize = 0;
        for (i, time) in self.times.iter().enumerate() {
            if time > &t {
                end_index = i;
                break;
            }
        }

        // Then send that to the integrator
        let t_start = times[end_index - 1];
        let t_end = times[end_index];
        self.integrator
            .interpolate(t_start, t_end, &self.dense[end_index - 1], t)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::callback::stop;
    use crate::controller::PIController;
    use crate::integrator::dormand_prince::DormandPrince45;
    use approx::assert_relative_eq;
    use nalgebra::Vector3;

    #[test]
    fn test_problem() {
        type Params = ();
        fn derivative(_t: f64, y: Vector3<f64>, _p: &Params) -> Vector3<f64> {
            y
        }
        let y0 = Vector3::new(1.0, 1.0, 1.0);

        let ode = ODE::new(&derivative, 0.0, 1.0, y0, ());
        let dp45 = DormandPrince45::new(1e-12_f64, 1e-5_f64);
        let controller = PIController::default();

        let mut problem = Problem::new(ode, dp45, controller);

        let solution = problem.solve();
        solution
            .times
            .iter()
            .zip(solution.states.iter())
            .for_each(|(time, state)| {
                assert_relative_eq!(state[0], time.exp(), max_relative = 1e-2);
            })
    }

    #[test]
    fn test_with_callback() {
        type Params = ();
        fn derivative(_t: f64, y: Vector3<f64>, _p: &Params) -> Vector3<f64> {
            y
        }
        let y0 = Vector3::new(1.0, 1.0, 1.0);

        let ode = ODE::new(&derivative, 0.0, 10.0, y0, ());
        let dp45 = DormandPrince45::new(1e-12_f64, 1e-5_f64);
        let controller = PIController::default();

        let value_too_high = Callback {
            event: &|_: f64, y: SVector<f64, 3>, _: &Params| 10.0 - y[0],
            effect: &stop,
        };

        let mut problem = Problem::new(ode, dp45, controller).with_callback(value_too_high);
        let solution = problem.solve();

        assert!(solution.states.last().unwrap()[0] == 10.0);
    }

    #[test]
    fn test_with_interpolation() {
        type Params = ();
        fn derivative(_t: f64, y: Vector3<f64>, _p: &Params) -> Vector3<f64> {
            y
        }
        let y0 = Vector3::new(1.0, 1.0, 1.0);

        let ode = ODE::new(&derivative, 0.0, 10.0, y0, ());
        let dp45 = DormandPrince45::new(1e-12_f64, 1e-6_f64);
        let controller = PIController::default();

        let mut problem = Problem::new(ode, dp45, controller);
        let solution = problem.solve();

        assert_relative_eq!(
            solution.interpolate(8.8)[0],
            8.8_f64.exp(),
            max_relative = 1e-6
        );
    }
}