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

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

// Simple 1D exponential decay problem
// y' = -k*y, y(0) = 1
fn bench_exponential_decay(c: &mut Criterion) {
    type Params = (f64,);
    let params = (0.5,);

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

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

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

    // Moderate tolerance - where BS3 should excel
    let tol = 1e-5;

    group.bench_function("bs3_tol_1e-5", |b| {
        let ode = ODE::new(&derivative, 0.0, 10.0, y0, params);
        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-5", |b| {
        let ode = ODE::new(&derivative, 0.0, 10.0, y0, params);
        let dp45 = DormandPrince45::new().a_tol(tol).r_tol(tol);
        b.iter(|| {
            black_box({
                Problem::new(ode, dp45, controller).solve();
            });
        });
    });

    group.finish();
}

// 2D harmonic oscillator
// y'' + y = 0, or as system: y1' = y2, y2' = -y1
fn bench_harmonic_oscillator(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");

    let tol = 1e-5;

    group.bench_function("bs3_tol_1e-5", |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-5", |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.finish();
}

// Nonlinear pendulum
// theta'' + (g/L)*sin(theta) = 0
fn bench_pendulum(c: &mut Criterion) {
    type Params = (f64, f64); // (g, L)
    let params = (9.81, 1.0);

    fn derivative(_t: f64, y: Vector2<f64>, p: &Params) -> Vector2<f64> {
        let &(g, l) = p;
        let theta = y[0];
        let d_theta = y[1];
        Vector2::new(d_theta, -(g / l) * theta.sin())
    }

    let y0 = Vector2::new(0.0, PI / 2.0); // Start from rest at angle 0, velocity PI/2
    let controller = PIController::default();

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

    let tol = 1e-6;

    group.bench_function("bs3_tol_1e-6", |b| {
        let ode = ODE::new(&derivative, 0.0, 10.0, y0, params);
        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-6", |b| {
        let ode = ODE::new(&derivative, 0.0, 10.0, y0, params);
        let dp45 = DormandPrince45::new().a_tol(tol).r_tol(tol);
        b.iter(|| {
            black_box({
                Problem::new(ode, dp45, controller).solve();
            });
        });
    });

    group.finish();
}

// 6D orbital mechanics - higher dimensional problem
fn bench_orbit_6d(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_6d");

    // Test at moderate tolerance
    let tol = 1e-6;

    group.bench_function("bs3_tol_1e-6", |b| {
        let ode = ODE::new(&derivative, 0.0, 10000.0, y0, params);
        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-6", |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.finish();
}

// Benchmark interpolation performance
fn bench_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("interpolation");

    let tol = 1e-6;

    // BS3 with interpolation
    group.bench_function("bs3_with_interpolation", |b| {
        let ode = ODE::new(&derivative, 0.0, 5.0, y0, ());
        let bs3 = BS3::new().a_tol(tol).r_tol(tol);
        b.iter(|| {
            black_box({
                let solution = Problem::new(ode, bs3, controller).solve();
                // Interpolate at 100 points
                let _: Vec<_> = (0..100).map(|i| solution.interpolate(i as f64 * 0.05)).collect();
            });
        });
    });

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

    group.finish();
}

// Tolerance scaling benchmark - how do methods perform at different tolerances?
fn bench_tolerance_scaling(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");

    let tolerances = [1e-3, 1e-4, 1e-5, 1e-6, 1e-7];

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

        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.finish();
}

criterion_group!(
    benches,
    bench_exponential_decay,
    bench_harmonic_oscillator,
    bench_pendulum,
    bench_orbit_6d,
    bench_interpolation,
    bench_tolerance_scaling,
);
criterion_main!(benches);