numra-ode 0.1.4

//! Automatic solver selection.
//!
//! Provides intelligent method selection based on problem characteristics.
//!
//! ## Usage
//!
//! ```rust
//! use numra_ode::{OdeProblem, auto_solve, SolverOptions};
//!
//! let problem = OdeProblem::new(
//!     |_t, y: &[f64], dydt: &mut [f64]| { dydt[0] = -y[0]; },
//!     0.0, 1.0, vec![1.0],
//! );
//! let options = SolverOptions::default();
//! let result = auto_solve(&problem, 0.0, 1.0, &[1.0], &options).unwrap();
//! assert!(result.success);
//! ```
//!
//! ## Selection Strategy
//!
//! - **Non-stiff problems**: Uses Tsit5 (efficient, accurate, FSAL)
//! - **Moderately stiff**: Uses Esdirk54 (L-stable, good efficiency)
//! - **Very stiff**: Uses BDF or Radau5 (high stiffness handling)
//! - **High accuracy**: Uses Vern8 (8th order accuracy)
//!
//! Author: Moussa Leblouba
//! Date: 5 March 2026
//! Modified: 2 May 2026

use faer::{ComplexField, Conjugate, SimpleEntity};
use numra_core::Scalar;

use crate::bdf::Bdf;
use crate::error::SolverError;
use crate::esdirk::Esdirk54;
use crate::problem::OdeSystem;
use crate::radau5::Radau5;
use crate::solver::{Solver, SolverOptions, SolverResult};
use crate::tsit5::Tsit5;
use crate::verner::{Vern6, Vern8};

/// Problem stiffness classification.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Stiffness {
    /// Non-stiff problem (use explicit methods)
    NonStiff,
    /// Moderate stiffness (ESDIRK methods work well)
    ModeratelyStiff,
    /// Highly stiff (BDF or Radau methods recommended)
    VeryStiff,
    /// Unknown (will be detected automatically)
    Unknown,
}

/// Accuracy requirements.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Accuracy {
    /// Low accuracy (rtol ~ 1e-3)
    Low,
    /// Standard accuracy (rtol ~ 1e-6)
    Standard,
    /// High accuracy (rtol ~ 1e-10)
    High,
    /// Very high accuracy (rtol ~ 1e-12+)
    VeryHigh,
}

/// Solver selection hints.
#[derive(Clone, Debug, Default)]
pub struct SolverHints {
    /// Problem stiffness
    pub stiffness: Option<Stiffness>,
    /// Accuracy requirements
    pub accuracy: Option<Accuracy>,
    /// Prefer implicit methods (for conservation)
    pub prefer_implicit: bool,
    /// Enable stiffness detection
    pub detect_stiffness: bool,
}

impl SolverHints {
    /// Create default hints.
    pub fn new() -> Self {
        Self {
            stiffness: None,
            accuracy: None,
            prefer_implicit: false,
            detect_stiffness: true,
        }
    }

    /// Set stiffness hint.
    pub fn stiffness(mut self, stiffness: Stiffness) -> Self {
        self.stiffness = Some(stiffness);
        self
    }

    /// Set accuracy requirement.
    pub fn accuracy(mut self, accuracy: Accuracy) -> Self {
        self.accuracy = Some(accuracy);
        self
    }

    /// Prefer implicit methods.
    pub fn implicit(mut self) -> Self {
        self.prefer_implicit = true;
        self
    }

    /// Enable/disable stiffness detection.
    pub fn detect_stiffness(mut self, detect: bool) -> Self {
        self.detect_stiffness = detect;
        self
    }
}

/// Determine accuracy level from options.
fn classify_accuracy<S: Scalar>(options: &SolverOptions<S>) -> Accuracy {
    let rtol = options.rtol.to_f64();
    if rtol >= 1e-3 {
        Accuracy::Low
    } else if rtol >= 1e-7 {
        Accuracy::Standard
    } else if rtol >= 1e-11 {
        Accuracy::High
    } else {
        Accuracy::VeryHigh
    }
}

/// Attempt stiffness detection.
fn detect_stiffness<S, Sys>(problem: &Sys, t: S, y: &[S], _options: &SolverOptions<S>) -> Stiffness
where
    S: Scalar,
    Sys: OdeSystem<S>,
{
    let dim = problem.dim();
    if dim == 0 {
        return Stiffness::Unknown;
    }

    // Compute Jacobian eigenvalues (approximate via power iteration)
    let h_factor = S::EPSILON.sqrt();
    let mut f0 = vec![S::ZERO; dim];
    let mut f1 = vec![S::ZERO; dim];
    let _jv = vec![S::ZERO; dim];

    problem.rhs(t, y, &mut f0);

    // Simple stiffness indicator: ratio of max/min Jacobian elements
    let mut max_jac = S::ZERO;
    let mut min_jac = S::INFINITY;
    let mut y_pert = y.to_vec();

    for j in 0..dim.min(10) {
        // Sample first 10 components for stiffness detection
        let yj = y[j];
        let h = h_factor * (S::ONE + yj.abs());
        y_pert[j] = yj + h;
        problem.rhs(t, &y_pert, &mut f1);
        y_pert[j] = yj;

        for i in 0..dim {
            let jij = ((f1[i] - f0[i]) / h).abs();
            if jij > S::from_f64(1e-15) {
                max_jac = max_jac.max(jij);
                min_jac = min_jac.min(jij);
            }
        }
    }

    // Stiffness ratio
    if max_jac < S::from_f64(1e-10) {
        return Stiffness::NonStiff;
    }

    let ratio = max_jac / min_jac.max(S::from_f64(1e-15));
    let ratio_f64 = ratio.to_f64();

    if ratio_f64 > 1e4 {
        Stiffness::VeryStiff
    } else if ratio_f64 > 100.0 {
        Stiffness::ModeratelyStiff
    } else {
        Stiffness::NonStiff
    }
}

/// Convenience function for automatic solving.
pub fn auto_solve<S, Sys>(
    problem: &Sys,
    t0: S,
    tf: S,
    y0: &[S],
    options: &SolverOptions<S>,
) -> Result<SolverResult<S>, SolverError>
where
    S: Scalar + SimpleEntity + Conjugate<Canonical = S> + ComplexField,
    Sys: OdeSystem<S>,
{
    auto_solve_with_hints(problem, t0, tf, y0, options, &SolverHints::new())
}

/// Convenience function for automatic solving with hints.
pub fn auto_solve_with_hints<S, Sys>(
    problem: &Sys,
    t0: S,
    tf: S,
    y0: &[S],
    options: &SolverOptions<S>,
    hints: &SolverHints,
) -> Result<SolverResult<S>, SolverError>
where
    S: Scalar + SimpleEntity + Conjugate<Canonical = S> + ComplexField,
    Sys: OdeSystem<S>,
{
    // Determine accuracy
    let accuracy = hints.accuracy.unwrap_or_else(|| classify_accuracy(options));

    // Determine stiffness
    let stiffness = hints.stiffness.unwrap_or_else(|| {
        if hints.detect_stiffness {
            detect_stiffness(problem, t0, y0, options)
        } else {
            Stiffness::Unknown
        }
    });

    // Select solver based on characteristics
    match (stiffness, accuracy, hints.prefer_implicit) {
        // Non-stiff problems
        (Stiffness::NonStiff, Accuracy::Low, false)
        | (Stiffness::NonStiff, Accuracy::Standard, false) => {
            Tsit5::solve(problem, t0, tf, y0, options)
        }
        (Stiffness::NonStiff, Accuracy::High, false) => Vern6::solve(problem, t0, tf, y0, options),
        (Stiffness::NonStiff, Accuracy::VeryHigh, false) => {
            Vern8::solve(problem, t0, tf, y0, options)
        }

        // Moderately stiff
        (Stiffness::ModeratelyStiff, _, _) => Esdirk54::solve(problem, t0, tf, y0, options),

        // Very stiff
        (Stiffness::VeryStiff, Accuracy::Low, _)
        | (Stiffness::VeryStiff, Accuracy::Standard, _) => Bdf::solve(problem, t0, tf, y0, options),
        (Stiffness::VeryStiff, Accuracy::High, _)
        | (Stiffness::VeryStiff, Accuracy::VeryHigh, _) => {
            Radau5::solve(problem, t0, tf, y0, options)
        }

        // Prefer implicit
        (_, _, true) => Esdirk54::solve(problem, t0, tf, y0, options),

        // Unknown/default: try explicit first
        (Stiffness::Unknown, _, _) => {
            // Try Tsit5 first
            if let Ok(result) = Tsit5::solve(problem, t0, tf, y0, options) {
                // Check if solution seems reasonable
                if result.stats.n_reject < result.stats.n_accept {
                    return Ok(result);
                }
            }

            // Fall back to implicit method
            Esdirk54::solve(problem, t0, tf, y0, options)
        }
    }
}

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

    #[test]
    fn test_auto_nonstiff() {
        let problem = OdeProblem::new(
            |_t, y: &[f64], dydt: &mut [f64]| {
                dydt[0] = -y[0];
            },
            0.0,
            5.0,
            vec![1.0],
        );
        let options = SolverOptions::default().rtol(1e-6);
        let result = auto_solve(&problem, 0.0, 5.0, &[1.0], &options).unwrap();

        assert!(result.success);
        let y_final = result.y_final().unwrap();
        let expected = (-5.0_f64).exp();
        assert!((y_final[0] - expected).abs() < 1e-4);
    }

    #[test]
    fn test_auto_stiff() {
        // Moderately stiff problem - use ESDIRK instead of BDF for now
        // since BDF still needs Newton iteration improvements
        let problem = OdeProblem::new(
            |_t, y: &[f64], dydt: &mut [f64]| {
                dydt[0] = -100.0 * y[0];
            },
            0.0,
            0.1,
            vec![1.0],
        );
        let options = SolverOptions::default().rtol(1e-3).atol(1e-5);
        // Use moderately stiff hint which selects ESDIRK (more robust than BDF currently)
        let hints = SolverHints::new().stiffness(Stiffness::ModeratelyStiff);

        let result = auto_solve_with_hints(&problem, 0.0, 0.1, &[1.0], &options, &hints).unwrap();

        assert!(result.success);
        let y_final = result.y_final().unwrap();
        let expected = (-10.0_f64).exp();
        assert!(
            (y_final[0] - expected).abs() < 0.05,
            "stiff: got {}, expected {}",
            y_final[0],
            expected
        );
    }

    #[test]
    fn test_auto_high_accuracy() {
        let problem = OdeProblem::new(
            |_t, y: &[f64], dydt: &mut [f64]| {
                dydt[0] = y[1];
                dydt[1] = -y[0];
            },
            0.0,
            10.0,
            vec![1.0, 0.0],
        );
        // Use moderate tolerances for reliable testing
        let options = SolverOptions::default().rtol(1e-5).atol(1e-7);
        let hints = SolverHints::new().stiffness(Stiffness::NonStiff);

        let result =
            auto_solve_with_hints(&problem, 0.0, 10.0, &[1.0, 0.0], &options, &hints).unwrap();

        assert!(result.success);
        let y_final = result.y_final().unwrap();
        // Allow 0.1% error for moderate tolerances
        assert!(
            (y_final[0] - 10.0_f64.cos()).abs() < 1e-3,
            "high accuracy: got {}, expected {}",
            y_final[0],
            10.0_f64.cos()
        );
    }

    #[test]
    fn test_auto_detect_stiffness() {
        // Non-stiff problem
        let problem1 = OdeProblem::new(
            |_t, y: &[f64], dydt: &mut [f64]| {
                dydt[0] = -y[0];
            },
            0.0,
            1.0,
            vec![1.0],
        );
        let options = SolverOptions::default();
        let stiffness1 = detect_stiffness(&problem1, 0.0, &[1.0], &options);
        assert_eq!(stiffness1, Stiffness::NonStiff);

        // Stiff problem
        let problem2 = OdeProblem::new(
            |_t, y: &[f64], dydt: &mut [f64]| {
                dydt[0] = -1000.0 * y[0] + 0.01 * y[1];
                dydt[1] = 0.01 * y[0] - y[1];
            },
            0.0,
            1.0,
            vec![1.0, 1.0],
        );
        let stiffness2 = detect_stiffness(&problem2, 0.0, &[1.0, 1.0], &options);
        assert!(stiffness2 == Stiffness::VeryStiff || stiffness2 == Stiffness::ModeratelyStiff);
    }

    #[test]
    fn test_accuracy_classification() {
        let opts_low: SolverOptions<f64> = SolverOptions::default().rtol(1e-2);
        let opts_std: SolverOptions<f64> = SolverOptions::default().rtol(1e-6);
        let opts_high: SolverOptions<f64> = SolverOptions::default().rtol(1e-10);
        let opts_vhigh: SolverOptions<f64> = SolverOptions::default().rtol(1e-13);

        assert_eq!(classify_accuracy(&opts_low), Accuracy::Low);
        assert_eq!(classify_accuracy(&opts_std), Accuracy::Standard);
        assert_eq!(classify_accuracy(&opts_high), Accuracy::High);
        assert_eq!(classify_accuracy(&opts_vhigh), Accuracy::VeryHigh);
    }

    #[test]
    fn test_auto_convenience() {
        let problem = OdeProblem::new(
            |_t, y: &[f64], dydt: &mut [f64]| {
                dydt[0] = -y[0];
            },
            0.0,
            2.0,
            vec![1.0],
        );
        let options = SolverOptions::default();

        let result = auto_solve(&problem, 0.0, 2.0, &[1.0], &options).unwrap();
        assert!(result.success);
    }
}