use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion};

use nalgebra::{Vector1, Vector2, Vector6};
use ordinary_diffeq::prelude::*;
use std::hint::black_box;

// Tight tolerance benchmarks - where Vern7 should excel
// Vern7 is designed for tolerances in the range 1e-8 to 1e-12

// Simple 1D exponential problem
// y' = y, y(0) = 1, solution: y(t) = e^t
fn bench_exponential_tight_tol(c: &mut Criterion) {
    type Params = ();

    fn derivative(_t: f64, y: Vector1<f64>, _p: &Params) -> Vector1<f64> {
        Vector1::new(y[0])
    }

    let y0 = Vector1::new(1.0);
    let controller = PIController::default();

    let mut group = c.benchmark_group("exponential_tight_tol");

    // Tight tolerance - where Vern7 should excel
    let tol = 1e-10;

    group.bench_function("bs3_tol_1e-10", |b| {
        let ode = ODE::new(&derivative, 0.0, 4.0, y0, ());
        let bs3 = BS3::new().a_tol(tol).r_tol(tol);
        b.iter(|| {
            black_box({
                Problem::new(ode, bs3, controller).solve();
            });
        });
    });

    group.bench_function("dp5_tol_1e-10", |b| {
        let ode = ODE::new(&derivative, 0.0, 4.0, y0, ());
        let dp45 = DormandPrince45::new().a_tol(tol).r_tol(tol);
        b.iter(|| {
            black_box({
                Problem::new(ode, dp45, controller).solve();
            });
        });
    });

    group.bench_function("vern7_tol_1e-10", |b| {
        let ode = ODE::new(&derivative, 0.0, 4.0, y0, ());
        let vern7 = Vern7::new().a_tol(tol).r_tol(tol);
        b.iter(|| {
            black_box({
                Problem::new(ode, vern7, controller).solve();
            });
        });
    });

    group.finish();
}

// 2D harmonic oscillator - smooth periodic system
// y'' + y = 0, or as system: y1' = y2, y2' = -y1
fn bench_harmonic_oscillator_tight_tol(c: &mut Criterion) {
    type Params = ();

    fn derivative(_t: f64, y: Vector2<f64>, _p: &Params) -> Vector2<f64> {
        Vector2::new(y[1], -y[0])
    }

    let y0 = Vector2::new(1.0, 0.0);
    let controller = PIController::default();

    let mut group = c.benchmark_group("harmonic_oscillator_tight_tol");

    let tol = 1e-10;

    group.bench_function("bs3_tol_1e-10", |b| {
        let ode = ODE::new(&derivative, 0.0, 20.0, y0, ());
        let bs3 = BS3::new().a_tol(tol).r_tol(tol);
        b.iter(|| {
            black_box({
                Problem::new(ode, bs3, controller).solve();
            });
        });
    });

    group.bench_function("dp5_tol_1e-10", |b| {
        let ode = ODE::new(&derivative, 0.0, 20.0, y0, ());
        let dp45 = DormandPrince45::new().a_tol(tol).r_tol(tol);
        b.iter(|| {
            black_box({
                Problem::new(ode, dp45, controller).solve();
            });
        });
    });

    group.bench_function("vern7_tol_1e-10", |b| {
        let ode = ODE::new(&derivative, 0.0, 20.0, y0, ());
        let vern7 = Vern7::new().a_tol(tol).r_tol(tol);
        b.iter(|| {
            black_box({
                Problem::new(ode, vern7, controller).solve();
            });
        });
    });

    group.finish();
}

// 6D orbital mechanics - high dimensional problem where tight tolerances matter
fn bench_orbit_tight_tol(c: &mut Criterion) {
    let mu = 3.98600441500000e14;

    type Params = (f64,);
    let params = (mu,);

    fn derivative(_t: f64, state: Vector6<f64>, p: &Params) -> Vector6<f64> {
        let acc = -(p.0 * state.fixed_rows::<3>(0)) / (state.fixed_rows::<3>(0).norm().powi(3));
        Vector6::new(state[3], state[4], state[5], acc[0], acc[1], acc[2])
    }

    let y0 = Vector6::new(
        4.263868426884883e6,
        5.146189057155391e6,
        1.1310208421331816e6,
        -5923.454461876975,
        4496.802639690076,
        1870.3893008991558,
    );

    let controller = PIController::new(0.37, 0.04, 10.0, 0.2, 1000.0, 0.9, 0.01);

    let mut group = c.benchmark_group("orbit_tight_tol");

    // Tight tolerance for orbital mechanics
    let tol = 1e-10;

    group.bench_function("dp5_tol_1e-10", |b| {
        let ode = ODE::new(&derivative, 0.0, 10000.0, y0, params);
        let dp45 = DormandPrince45::new().a_tol(tol).r_tol(tol);
        b.iter(|| {
            black_box({
                Problem::new(ode, dp45, controller).solve();
            });
        });
    });

    group.bench_function("vern7_tol_1e-10", |b| {
        let ode = ODE::new(&derivative, 0.0, 10000.0, y0, params);
        let vern7 = Vern7::new().a_tol(tol).r_tol(tol);
        b.iter(|| {
            black_box({
                Problem::new(ode, vern7, controller).solve();
            });
        });
    });

    group.finish();
}

// Benchmark interpolation performance with lazy dense output
fn bench_vern7_interpolation(c: &mut Criterion) {
    type Params = ();

    fn derivative(_t: f64, y: Vector1<f64>, _p: &Params) -> Vector1<f64> {
        Vector1::new(y[0])
    }

    let y0 = Vector1::new(1.0);
    let controller = PIController::default();

    let mut group = c.benchmark_group("vern7_interpolation");

    let tol = 1e-10;

    // Vern7 with interpolation (should compute extra stages lazily)
    group.bench_function("vern7_with_interpolation", |b| {
        b.iter(|| {
            black_box({
                let ode = ODE::new(&derivative, 0.0, 5.0, y0, ());
                let vern7 = Vern7::new().a_tol(tol).r_tol(tol);
                let mut problem = Problem::new(ode, vern7, controller);
                let solution = problem.solve();
                // Interpolate at 100 points - first one computes extra stages
                let _: Vec<_> = (0..100).map(|i| solution.interpolate(i as f64 * 0.05)).collect();
            });
        });
    });

    // DP5 with interpolation for comparison
    group.bench_function("dp5_with_interpolation", |b| {
        b.iter(|| {
            black_box({
                let ode = ODE::new(&derivative, 0.0, 5.0, y0, ());
                let dp45 = DormandPrince45::new().a_tol(tol).r_tol(tol);
                let mut problem = Problem::new(ode, dp45, controller);
                let solution = problem.solve();
                let _: Vec<_> = (0..100).map(|i| solution.interpolate(i as f64 * 0.05)).collect();
            });
        });
    });

    group.finish();
}

// Tolerance scaling for Vern7 vs lower-order methods
fn bench_tolerance_scaling_vern7(c: &mut Criterion) {
    type Params = ();

    fn derivative(_t: f64, y: Vector1<f64>, _p: &Params) -> Vector1<f64> {
        Vector1::new(-y[0])
    }

    let y0 = Vector1::new(1.0);
    let controller = PIController::default();

    let mut group = c.benchmark_group("tolerance_scaling_vern7");

    // Focus on tight tolerances where Vern7 excels
    let tolerances = [1e-6, 1e-7, 1e-8, 1e-9, 1e-10];

    for &tol in &tolerances {
        group.bench_with_input(BenchmarkId::new("dp5", tol), &tol, |b, &tol| {
            let ode = ODE::new(&derivative, 0.0, 10.0, y0, ());
            let dp45 = DormandPrince45::new().a_tol(tol).r_tol(tol);
            b.iter(|| {
                black_box({
                    Problem::new(ode, dp45, controller).solve();
                });
            });
        });

        group.bench_with_input(BenchmarkId::new("vern7", tol), &tol, |b, &tol| {
            let ode = ODE::new(&derivative, 0.0, 10.0, y0, ());
            let vern7 = Vern7::new().a_tol(tol).r_tol(tol);
            b.iter(|| {
                black_box({
                    Problem::new(ode, vern7, controller).solve();
                });
            });
        });
    }

    group.finish();
}

criterion_group!(
    benches,
    bench_exponential_tight_tol,
    bench_harmonic_oscillator_tight_tol,
    bench_orbit_tight_tol,
    bench_vern7_interpolation,
    bench_tolerance_scaling_vern7,
);
criterion_main!(benches);