conspire 0.6.0

use crate::math::{
    Banded, Scalar, Tensor, TensorVec, Vector,
    integrate::{
        ImplicitDaeFirstOrderMinimize, ImplicitDaeFirstOrderRoot, ImplicitDaeSecondOrderMinimize,
        ImplicitDaeZerothOrderRoot, IntegrationError, VariableStepExplicit,
    },
    optimize::{
        EqualityConstraint, FirstOrderOptimization, FirstOrderRootFinding, SecondOrderOptimization,
        ZerothOrderRootFinding,
    },
};
use std::ops::{Mul, Sub};

/// Variable-step explicit integrators for implicit differential-algebraic equations.
pub trait ImplicitDaeVariableStepExplicit<Y, U>
where
    Self: VariableStepExplicit<Y, U>,
    Y: Tensor,
    U: TensorVec<Item = Y>,
    for<'a> &'a Y: Mul<Scalar, Output = Y> + Sub<&'a Y, Output = Y>,
{
    fn integrate_implicit_dae_variable_step(
        &self,
        mut evolution: impl FnMut(Scalar, &Y, &Y) -> Result<Y, String>,
        time: &[Scalar],
        initial_condition: Y,
    ) -> Result<(Vector, U, U), IntegrationError> {
        let t_0 = time[0];
        let t_f = time[time.len() - 1];
        if time.len() < 2 {
            return Err(IntegrationError::LengthTimeLessThanTwo);
        } else if t_0 >= t_f {
            return Err(IntegrationError::InitialTimeNotLessThanFinalTime);
        }
        let mut t = t_0;
        let mut dt = t_f - t_0;
        let mut t_sol = Vector::new();
        t_sol.push(t_0);
        let mut dydt = &initial_condition * 0.0;
        let mut y = initial_condition;
        let mut k = vec![Y::default(); Self::SLOPES];
        k[0] = evolution(t, &y, &dydt)?;
        let mut y_sol = U::new();
        y_sol.push(y.clone());
        let mut dydt_sol = U::new();
        dydt_sol.push(k[0].clone());
        let mut y_trial = Y::default();
        while t < t_f {
            match self.slopes_and_error(
                |t: Scalar, y: &Y| evolution(t, y, &dydt),
                &y,
                t,
                dt,
                &mut k,
                &mut y_trial,
            ) {
                Ok(e) => {
                    if let Some(error) = self
                        .step(
                            |t: Scalar, y: &Y| evolution(t, y, &dydt),
                            &mut y,
                            &mut t,
                            &mut y_sol,
                            &mut t_sol,
                            &mut dydt_sol,
                            &mut dt,
                            &mut k,
                            &y_trial,
                            e,
                        )
                        .err()
                    {
                        dt *= self.dt_cut();
                        if dt < self.dt_min() {
                            return Err(IntegrationError::MinimumStepSizeUpstream(
                                self.dt_min(),
                                error,
                                format!("{:?}", self),
                            ));
                        }
                    } else {
                        dydt = k[0].clone();
                        dt = dt.min(t_f - t);
                        if dt < self.dt_min() && t < t_f {
                            return Err(IntegrationError::MinimumStepSizeReached(
                                self.dt_min(),
                                format!("{:?}", self),
                            ));
                        }
                    }
                }
                Err(error) => {
                    dt *= self.dt_cut();
                    if dt < self.dt_min() {
                        return Err(IntegrationError::MinimumStepSizeUpstream(
                            self.dt_min(),
                            error,
                            format!("{:?}", self),
                        ));
                    }
                }
            }
        }
        if time.len() > 2 {
            let t_int = Vector::from(time);
            let (y_int, dydt_int) = self.interpolate_implicit_dae_variable_step(
                evolution, &t_int, &t_sol, &y_sol, &dydt_sol,
            )?;
            Ok((t_int, y_int, dydt_int))
        } else {
            Ok((t_sol, y_sol, dydt_sol))
        }
    }
    fn interpolate_implicit_dae_variable_step(
        &self,
        mut evolution: impl FnMut(Scalar, &Y, &Y) -> Result<Y, String>,
        time: &Vector,
        tp: &Vector,
        yp: &U,
        dydtp: &U,
    ) -> Result<(U, U), IntegrationError> {
        let mut dt;
        let mut i;
        let mut k = vec![Y::default(); Self::SLOPES];
        let mut t;
        let mut y;
        let mut dydt;
        let mut y_int = U::new();
        let mut dydt_int = U::new();
        let mut y_trial = Y::default();
        for time_k in time.iter() {
            i = tp.iter().position(|tp_i| tp_i >= time_k).unwrap();
            if time_k == &tp[i] {
                t = tp[i];
                y_trial = yp[i].clone();
                dt = 0.0;
            } else {
                t = tp[i - 1];
                y = &yp[i - 1];
                dydt = &dydtp[i - 1];
                dt = time_k - t;
                k[0] = evolution(t, y, dydt)?;
                Self::slopes(
                    |t: Scalar, y: &Y| evolution(t, y, dydt),
                    y,
                    t,
                    dt,
                    &mut k,
                    &mut y_trial,
                )?;
            }
            dydt_int.push(evolution(t + dt, &y_trial, &k[0])?);
            y_int.push(y_trial.clone());
        }
        Ok((y_int, dydt_int))
    }
}

impl<I, Y, U> ImplicitDaeVariableStepExplicit<Y, U> for I
where
    Self: VariableStepExplicit<Y, U>,
    Y: Tensor,
    U: TensorVec<Item = Y>,
    for<'a> &'a Y: Mul<Scalar, Output = Y> + Sub<&'a Y, Output = Y>,
{
}

/// Variable-step explicit integrators for implicit differential-algebraic equations using zeroth-order root-finding.
pub trait ImplicitDaeVariableStepExplicitZerothOrderRoot<Y, U>
where
    Self: ImplicitDaeVariableStepExplicit<Y, U>,
    Y: Tensor,
    U: TensorVec<Item = Y>,
    for<'a> &'a Y: Mul<Scalar, Output = Y> + Sub<&'a Y, Output = Y>,
{
    fn integrate_implicit_dae_variable_step_explicit_root_0(
        &self,
        mut function: impl FnMut(Scalar, &Y, &Y) -> Result<Y, String>,
        solver: impl ZerothOrderRootFinding<Y>,
        time: &[Scalar],
        initial_condition: Y,
        mut equality_constraint: impl FnMut(Scalar) -> EqualityConstraint,
    ) -> Result<(Vector, U, U), IntegrationError> {
        let evolution = |t: Scalar, y: &Y, dydt_0: &Y| -> Result<Y, String> {
            Ok(solver.root(
                |dydt| function(t, y, dydt),
                dydt_0.clone(),
                equality_constraint(t),
            )?)
        };
        self.integrate_implicit_dae_variable_step(evolution, time, initial_condition)
    }
}

impl<I, Y, U> ImplicitDaeVariableStepExplicitZerothOrderRoot<Y, U> for I
where
    I: ImplicitDaeVariableStepExplicit<Y, U>,
    Y: Tensor,
    U: TensorVec<Item = Y>,
    for<'a> &'a Y: Mul<Scalar, Output = Y> + Sub<&'a Y, Output = Y>,
{
}

impl<I, Y, U> ImplicitDaeZerothOrderRoot<Y, U> for I
where
    Self: ImplicitDaeVariableStepExplicitZerothOrderRoot<Y, U>,
    Y: Tensor,
    U: TensorVec<Item = Y>,
    for<'a> &'a Y: Mul<Scalar, Output = Y> + Sub<&'a Y, Output = Y>,
{
    fn integrate(
        &self,
        function: impl FnMut(Scalar, &Y, &Y) -> Result<Y, String>,
        solver: impl ZerothOrderRootFinding<Y>,
        time: &[Scalar],
        initial_condition: Y,
        equality_constraint: impl FnMut(Scalar) -> EqualityConstraint,
    ) -> Result<(Vector, U, U), IntegrationError> {
        self.integrate_implicit_dae_variable_step_explicit_root_0(
            function,
            solver,
            time,
            initial_condition,
            equality_constraint,
        )
    }
}

/// Variable-step explicit integrators for implicit differential-algebraic equations using first-order root-finding.
pub trait ImplicitDaeVariableStepExplicitFirstOrderRoot<F, J, Y, U>
where
    Self: ImplicitDaeVariableStepExplicit<Y, U>,
    Y: Tensor,
    U: TensorVec<Item = Y>,
    for<'a> &'a Y: Mul<Scalar, Output = Y> + Sub<&'a Y, Output = Y>,
{
    fn integrate_implicit_dae_variable_step_explicit_root_1(
        &self,
        mut function: impl FnMut(Scalar, &Y, &Y) -> Result<F, String>,
        mut jacobian: impl FnMut(Scalar, &Y, &Y) -> Result<J, String>,
        solver: impl FirstOrderRootFinding<F, J, Y>,
        time: &[Scalar],
        initial_condition: Y,
        mut equality_constraint: impl FnMut(Scalar) -> EqualityConstraint,
    ) -> Result<(Vector, U, U), IntegrationError> {
        let evolution = |t: Scalar, y: &Y, dydt_0: &Y| -> Result<Y, String> {
            Ok(solver.root(
                |dydt| function(t, y, dydt),
                |dydt| jacobian(t, y, dydt),
                dydt_0.clone(),
                equality_constraint(t),
            )?)
        };
        self.integrate_implicit_dae_variable_step(evolution, time, initial_condition)
    }
}

impl<I, F, J, Y, U> ImplicitDaeVariableStepExplicitFirstOrderRoot<F, J, Y, U> for I
where
    I: ImplicitDaeVariableStepExplicit<Y, U>,
    Y: Tensor,
    U: TensorVec<Item = Y>,
    for<'a> &'a Y: Mul<Scalar, Output = Y> + Sub<&'a Y, Output = Y>,
{
}

impl<I, F, J, Y, U> ImplicitDaeFirstOrderRoot<F, J, Y, U> for I
where
    Self: ImplicitDaeVariableStepExplicitFirstOrderRoot<F, J, Y, U>,
    Y: Tensor,
    U: TensorVec<Item = Y>,
    for<'a> &'a Y: Mul<Scalar, Output = Y> + Sub<&'a Y, Output = Y>,
{
    fn integrate(
        &self,
        function: impl FnMut(Scalar, &Y, &Y) -> Result<F, String>,
        jacobian: impl FnMut(Scalar, &Y, &Y) -> Result<J, String>,
        solver: impl FirstOrderRootFinding<F, J, Y>,
        time: &[Scalar],
        initial_condition: Y,
        equality_constraint: impl FnMut(Scalar) -> EqualityConstraint,
    ) -> Result<(Vector, U, U), IntegrationError> {
        self.integrate_implicit_dae_variable_step_explicit_root_1(
            function,
            jacobian,
            solver,
            time,
            initial_condition,
            equality_constraint,
        )
    }
}

/// Variable-step explicit integrators for implicit differential-algebraic equations using first-order minimization.
pub trait ImplicitDaeVariableStepExplicitFirstOrderMinimize<F, Y, U>
where
    Self: ImplicitDaeVariableStepExplicit<Y, U>,
    Y: Tensor,
    U: TensorVec<Item = Y>,
    for<'a> &'a Y: Mul<Scalar, Output = Y> + Sub<&'a Y, Output = Y>,
{
    #[allow(clippy::too_many_arguments)]
    fn integrate_implicit_dae_variable_step_explicit_minimize_1(
        &self,
        mut function: impl FnMut(Scalar, &Y, &Y) -> Result<F, String>,
        mut jacobian: impl FnMut(Scalar, &Y, &Y) -> Result<Y, String>,
        solver: impl FirstOrderOptimization<F, Y>,
        time: &[Scalar],
        initial_condition: Y,
        mut equality_constraint: impl FnMut(Scalar) -> EqualityConstraint,
    ) -> Result<(Vector, U, U), IntegrationError> {
        let evolution = |t: Scalar, y: &Y, dydt_0: &Y| -> Result<Y, String> {
            Ok(solver.minimize(
                |dydt| function(t, y, dydt),
                |dydt| jacobian(t, y, dydt),
                dydt_0.clone(),
                equality_constraint(t),
            )?)
        };
        self.integrate_implicit_dae_variable_step(evolution, time, initial_condition)
    }
}

impl<I, F, Y, U> ImplicitDaeVariableStepExplicitFirstOrderMinimize<F, Y, U> for I
where
    I: ImplicitDaeVariableStepExplicit<Y, U>,
    Y: Tensor,
    U: TensorVec<Item = Y>,
    for<'a> &'a Y: Mul<Scalar, Output = Y> + Sub<&'a Y, Output = Y>,
{
}

impl<I, F, Y, U> ImplicitDaeFirstOrderMinimize<F, Y, U> for I
where
    Self: ImplicitDaeVariableStepExplicitFirstOrderMinimize<F, Y, U>,
    Y: Tensor,
    U: TensorVec<Item = Y>,
    for<'a> &'a Y: Mul<Scalar, Output = Y> + Sub<&'a Y, Output = Y>,
{
    fn integrate(
        &self,
        function: impl FnMut(Scalar, &Y, &Y) -> Result<F, String>,
        jacobian: impl FnMut(Scalar, &Y, &Y) -> Result<Y, String>,
        solver: impl FirstOrderOptimization<F, Y>,
        time: &[Scalar],
        initial_condition: Y,
        equality_constraint: impl FnMut(Scalar) -> EqualityConstraint,
    ) -> Result<(Vector, U, U), IntegrationError> {
        self.integrate_implicit_dae_variable_step_explicit_minimize_1(
            function,
            jacobian,
            solver,
            time,
            initial_condition,
            equality_constraint,
        )
    }
}

/// Variable-step explicit integrators for implicit differential-algebraic equations using second-order minimization.
pub trait ImplicitDaeVariableStepExplicitSecondOrderMinimize<F, J, H, Y, U>
where
    Self: ImplicitDaeVariableStepExplicit<Y, U>,
    Y: Tensor,
    U: TensorVec<Item = Y>,
    for<'a> &'a Y: Mul<Scalar, Output = Y> + Sub<&'a Y, Output = Y>,
{
    #[allow(clippy::too_many_arguments)]
    fn integrate_implicit_dae_variable_step_explicit_minimize_2(
        &self,
        mut function: impl FnMut(Scalar, &Y, &Y) -> Result<F, String>,
        mut jacobian: impl FnMut(Scalar, &Y, &Y) -> Result<J, String>,
        mut hessian: impl FnMut(Scalar, &Y, &Y) -> Result<H, String>,
        solver: impl SecondOrderOptimization<F, J, H, Y>,
        time: &[Scalar],
        initial_condition: Y,
        mut equality_constraint: impl FnMut(Scalar) -> EqualityConstraint,
        banded: Option<Banded>,
    ) -> Result<(Vector, U, U), IntegrationError> {
        let evolution = |t: Scalar, y: &Y, dydt_0: &Y| -> Result<Y, String> {
            Ok(solver.minimize(
                |dydt| function(t, y, dydt),
                |dydt| jacobian(t, y, dydt),
                |dydt| hessian(t, y, dydt),
                dydt_0.clone(),
                equality_constraint(t),
                banded.clone(),
            )?)
        };
        self.integrate_implicit_dae_variable_step(evolution, time, initial_condition)
    }
}

impl<I, F, J, H, Y, U> ImplicitDaeVariableStepExplicitSecondOrderMinimize<F, J, H, Y, U> for I
where
    I: ImplicitDaeVariableStepExplicit<Y, U>,
    Y: Tensor,
    U: TensorVec<Item = Y>,
    for<'a> &'a Y: Mul<Scalar, Output = Y> + Sub<&'a Y, Output = Y>,
{
}

impl<I, F, J, H, Y, U> ImplicitDaeSecondOrderMinimize<F, J, H, Y, U> for I
where
    Self: ImplicitDaeVariableStepExplicitSecondOrderMinimize<F, J, H, Y, U>,
    Y: Tensor,
    U: TensorVec<Item = Y>,
    for<'a> &'a Y: Mul<Scalar, Output = Y> + Sub<&'a Y, Output = Y>,
{
    fn integrate(
        &self,
        function: impl FnMut(Scalar, &Y, &Y) -> Result<F, String>,
        jacobian: impl FnMut(Scalar, &Y, &Y) -> Result<J, String>,
        hessian: impl FnMut(Scalar, &Y, &Y) -> Result<H, String>,
        solver: impl SecondOrderOptimization<F, J, H, Y>,
        time: &[Scalar],
        initial_condition: Y,
        equality_constraint: impl FnMut(Scalar) -> EqualityConstraint,
        banded: Option<Banded>,
    ) -> Result<(Vector, U, U), IntegrationError> {
        self.integrate_implicit_dae_variable_step_explicit_minimize_2(
            function,
            jacobian,
            hessian,
            solver,
            time,
            initial_condition,
            equality_constraint,
            banded,
        )
    }
}