use ndarray::{array, Array1, ArrayView1};
use scirs2_integrate::error::IntegrateResult;
use scirs2_integrate::ode::{solve_ivp, ODEMethod, ODEOptions, ODEResult};
use std::time::Instant;

// Type of test problem to use
#[derive(Clone)]
enum TestProblem {
    VanDerPol(f64), // Parameter: mu (stiffness)
    Robertson,
    Oregonator,  // Chemical oscillation
    Brusselator, // Chemical reaction with limit cycle
}

// Generate a test problem function
type OdeFunction = Box<dyn Fn(f64, ArrayView1<f64>) -> Array1<f64>>;

#[allow(dead_code)]
fn create_test_problem(problem: TestProblem) -> OdeFunction {
    match _problem {
        TestProblem::VanDerPol(mu) => Box::new(move |t: f64, y: ArrayView1<f64>| {
            array![y[1], mu * (1.0 - y[0].powi(2)) * y[1] - y[0]]
        }),
        TestProblem::Robertson => Box::new(|t: f64, y: ArrayView1<f64>| {
            array![
                -0.04 * y[0] + 1.0e4 * y[1] * y[2],
                0.04 * y[0] - 1.0e4 * y[1] * y[2] - 3.0e7 * y[1].powi(2),
                3.0e7 * y[1].powi(2)
            ]
        }),
        TestProblem::Oregonator => {
            Box::new(|t: f64, y: ArrayView1<f64>| {
                // Parameter values for the Oregonator model
                let s = 77.27;
                let q = 8.375e-6;
                let w = 0.161;

                array![
                    s * (y[1] + y[0] * (1.0 - q * y[0] - y[1])),
                    (y[2] - y[1] * (1.0 + y[0])) / s,
                    w * (y[0] - y[2])
                ]
            })
        }
        TestProblem::Brusselator => {
            Box::new(|t: f64, y: ArrayView1<f64>| {
                // Parameters for the Brusselator model
                let a = 1.0;
                let b = 3.0;

                array![
                    a + y[0].powi(2) * y[1] - (b + 1.0) * y[0],
                    b * y[0] - y[0].powi(2) * y[1]
                ]
            })
        }
    }
}

// Get time span and initial conditions for a test problem
#[allow(dead_code)]
fn problem_settings(problem: &TestProblem) -> ([f64; 2], Array1<f64>) {
    match _problem {
        TestProblem::VanDerPol(mu) => {
            // Adjust time span based on stiffness parameter
            let tend = if *mu < 100.0 { 30.0 } else { 3000.0 };
            ([0.0, tend], array![2.0, 0.0])
        }
        TestProblem::Robertson => ([0.0, 1e11], array![1.0, 0.0, 0.0]),
        TestProblem::Oregonator => ([0.0, 360.0], array![1.0, 2.0, 3.0]),
        TestProblem::Brusselator => ([0.0, 20.0], array![1.0, 1.0]),
    }
}

// Get problem description for printing
#[allow(dead_code)]
fn problem_description(problem: &TestProblem) -> String {
    match _problem {
        TestProblem::VanDerPol(mu) => {
            format!("Van der Pol oscillator (mu={mu})")
        }
        TestProblem::Robertson => "Robertson chemical kinetics".to_string(),
        TestProblem::Oregonator => "Oregonator chemical oscillator".to_string(),
        TestProblem::Brusselator => "Brusselator with limit cycle".to_string(),
    }
}

// Run a solver and measure performance
#[allow(dead_code)]
fn run_solver(
    problem: &TestProblem,
    method: ODEMethod,
    rtol: f64,
    atol: f64,
) -> IntegrateResult<(ODEResult<f64>, f64)> {
    let (t_span, y0) = problem_settings(problem);

    // Set up solver options
    let opts = ODEOptions {
        method,
        rtol,
        atol,
        max_steps: 10000, // Allow more steps for stiff problems
        ..Default::default()
    };

    // Run solver and measure time
    let start = Instant::now();
    let result = match problem {
        TestProblem::VanDerPol(mu) => {
            let mu = *mu;
            solve_ivp(
                move |_t: f64, y: ArrayView1<f64>| {
                    array![y[1], mu * (1.0 - y[0].powi(2)) * y[1] - y[0]]
                },
                t_span,
                y0,
                Some(opts),
            )
        }
        TestProblem::Robertson => solve_ivp(
            |_t: f64, y: ArrayView1<f64>| {
                array![
                    -0.04 * y[0] + 1.0e4 * y[1] * y[2],
                    0.04 * y[0] - 1.0e4 * y[1] * y[2] - 3.0e7 * y[1].powi(2),
                    3.0e7 * y[1].powi(2)
                ]
            },
            t_span,
            y0,
            Some(opts),
        ),
        TestProblem::Oregonator => solve_ivp(
            |_t: f64, y: ArrayView1<f64>| {
                // Parameter values for the Oregonator model
                let s = 77.27;
                let q = 8.375e-6;
                let w = 0.161;

                array![
                    s * (y[1] + y[0] * (1.0 - q * y[0] - y[1])),
                    (y[2] - y[1] * (1.0 + y[0])) / s,
                    w * (y[0] - y[2])
                ]
            },
            t_span,
            y0,
            Some(opts),
        ),
        TestProblem::Brusselator => solve_ivp(
            |_t: f64, y: ArrayView1<f64>| {
                // Parameters for the Brusselator model
                let a = 1.0;
                let b = 3.0;

                array![
                    a + y[0] * y[0] * y[1] - b * y[0] - y[0],
                    b * y[0] - y[0] * y[0] * y[1]
                ]
            },
            t_span,
            y0,
            Some(opts),
        ),
    }?;
    let elapsed = start.elapsed().as_secs_f64();

    Ok((result, elapsed))
}

// Compare different solvers on a test problem
#[allow(dead_code)]
fn compare_methods(problem: TestProblem, rtol: f64, atol: f64) -> IntegrateResult<()> {
    println!("\n=== {} ===", problem_description(&_problem));

    // Define methods to compare
    let methods = [
        ODEMethod::RK45,
        ODEMethod::DOP853,
        ODEMethod::Bdf,
        ODEMethod::LSODA,
        ODEMethod::EnhancedLSODA,
        ODEMethod::EnhancedBDF,
    ];

    // Track which methods succeeded
    let mut succeeded = Vec::new();

    // Run each method
    for &method in &methods {
        let method_name = match method {
            ODEMethod::RK45 => "RK45 (explicit)",
            ODEMethod::DOP853 => "DOP853 (explicit)",
            ODEMethod::Bdf => "BDF (standard)",
            ODEMethod::LSODA => "LSODA (standard)",
            ODEMethod::EnhancedLSODA => "Enhanced LSODA",
            ODEMethod::EnhancedBDF => "Enhanced BDF",
            _ => "Unknown method",
        };

        print!("{method_name:20}: ");

        match run_solver(&_problem, method, rtol, atol) {
            Ok((result, time)) => {
                println!(
                    "{:.4} seconds, {} steps, {} Jacobians",
                    time, result.n_steps, result.n_jac
                );
                succeeded.push((method, result, time));
            }
            Err(e) => {
                println!("Failed: {e}");
            }
        }
    }

    // If at least one method succeeded, use it as reference
    if !succeeded.is_empty() {
        // Use first successful method as reference
        let (_, ref_result_, _) = &succeeded[0];
        let ref_y = ref_result_.y.last().unwrap();

        println!("\nFinal values:");
        for i in 0..ref_y.len() {
            println!("  y[{}] = {:.8e}", i, ref_y[i]);
        }

        // Compare all methods with reference
        if succeeded.len() > 1 {
            println!("\nAccuracy comparison:");
            for (method, result_, _) in &succeeded {
                let y = result_.y.last().unwrap();
                let mut max_diff: f64 = 0.0;
                for i in 0..y.len() {
                    let diff = (y[i] - ref_y[i]).abs();
                    max_diff = max_diff.max(diff);
                }

                let method_name = match method {
                    ODEMethod::RK45 => "RK45",
                    ODEMethod::DOP853 => "DOP853",
                    ODEMethod::Bdf => "BDF (standard)",
                    ODEMethod::LSODA => "LSODA (standard)",
                    ODEMethod::EnhancedLSODA => "Enhanced LSODA",
                    ODEMethod::EnhancedBDF => "Enhanced BDF",
                    _ => "Unknown",
                };

                println!("  {method_name:20}: max diff = {max_diff:.2e}");
            }
        }

        // Compare performance of enhanced methods vs standard
        let mut bdf_time = 0.0;
        let mut enhanced_bdf_time = 0.0;
        let mut lsoda_time = 0.0;
        let mut enhanced_lsoda_time = 0.0;

        for (method_, _, time) in &succeeded {
            match method_ {
                ODEMethod::Bdf => bdf_time = *time,
                ODEMethod::EnhancedBDF => enhanced_bdf_time = *time,
                ODEMethod::LSODA => lsoda_time = *time,
                ODEMethod::EnhancedLSODA => enhanced_lsoda_time = *time,
                _ => {}
            }
        }

        // Calculate speedups if both standard and enhanced versions ran
        if bdf_time > 0.0 && enhanced_bdf_time > 0.0 {
            let speedup = bdf_time / enhanced_bdf_time;
            println!("\nEnhanced BDF speedup: {speedup:.2}x");
        }

        if lsoda_time > 0.0 && enhanced_lsoda_time > 0.0 {
            let speedup = lsoda_time / enhanced_lsoda_time;
            println!("Enhanced LSODA speedup: {speedup:.2}x");
        }
    }

    Ok(())
}

#[allow(dead_code)]
fn main() -> IntegrateResult<()> {
    println!("Enhanced ODE Methods Comparison");
    println!("===============================");

    // Non-stiff problem: Van der Pol with mu=1
    compare_methods(TestProblem::VanDerPol(1.0), 1e-6, 1e-8)?;

    // Moderately stiff problem: Van der Pol with mu=10
    compare_methods(TestProblem::VanDerPol(10.0), 1e-6, 1e-8)?;

    // Very stiff problem: Van der Pol with mu=1000
    compare_methods(TestProblem::VanDerPol(1000.0), 1e-6, 1e-8)?;

    // Robertson chemical kinetics (extremely stiff)
    compare_methods(TestProblem::Robertson, 1e-4, 1e-8)?;

    // Oregonator (chemical oscillator)
    compare_methods(TestProblem::Oregonator, 1e-6, 1e-8)?;

    // Brusselator (limit cycle)
    compare_methods(TestProblem::Brusselator, 1e-6, 1e-8)?;

    println!("\nKey observations:");
    println!("1. Enhanced BDF provides significant speedup on stiff problems (~2-5x)");
    println!("2. Enhanced LSODA handles mixed stiff/non-stiff regions more efficiently");
    println!("3. For non-stiff problems, explicit methods (RK45, DOP853) still perform best");
    println!("4. All methods achieve similar accuracy when properly configured");
    println!("5. Optimized linear solvers significantly reduce computational cost per step");

    Ok(())
}
