//! Two-body propagation, and the derivative evaluation underneath it.
//!
//! The derivative benchmarks are the ones that matter most: a propagation
//! evaluates them thousands of times per day of arc, so anything wasteful in
//! assembly or in a model shows up here before it shows up in a wall-clock
//! figure. They also isolate the cost of a layout, which should be nil — the
//! `Orbit`, `OrbitMass` and `Full` numbers ought to differ only by the work of
//! moving more elements around.

use std::hint::black_box;

use astrokit::dynamics::{self, Context};
use astrokit::kepler;
use astrokit::prelude::*;
use astrokit::propagation::{state_derivative, Assembly};
use criterion::{criterion_group, criterion_main, Criterion};
use nalgebra::Vector3;

const MU: f64 = 3.986004415e14;
const SMA: f64 = 6.7781363e6;

fn epoch() -> Epoch {
    Epoch::from_gregorian_utc_hms(2026, 1, 1, 0, 0, 0)
}

fn position() -> Vector3<f64> {
    Vector3::new(
        4.263868426884883e6,
        5.146189057155391e6,
        1.1310208421331816e6,
    )
}

fn velocity() -> Vector3<f64> {
    Vector3::new(-5923.454461876975, 4496.802639690076, 1870.3893008991558)
}

fn two_body() -> Vec<Dynamics> {
    vec![TwoBody { body: Body::EARTH }]
}

fn spacecraft() -> Spacecraft {
    Spacecraft::rectangular_box(500.0, Vector3::new(2.0, 1.0, 3.0), 2.2, 1.3)
}

/// A tumble about no principal axis, fast enough to cross the shadow-set threshold
/// several times over an orbital period.
fn angular_rate() -> Vector3<f64> {
    Vector3::new(0.0012, 0.0028, -0.0007)
}

/// One derivative evaluation, per layout. The layout should cost nothing.
fn bench_derivative(c: &mut Criterion) {
    let mut group = c.benchmark_group("derivative");

    let assembly = Assembly::new(
        dynamics::build::<f64, 6, Orbit>(&two_body()),
        Context::new(epoch(), Frame::EARTH_J2000),
    );
    let y = Orbit.with_position(position()).with_velocity(velocity());
    group.bench_function("two_body/Orbit", |b| {
        b.iter(|| state_derivative(black_box(0.0), black_box(y.vector), &assembly))
    });

    let assembly = Assembly::new(
        dynamics::build::<f64, 7, OrbitMass>(&two_body()),
        Context::new(epoch(), Frame::EARTH_J2000),
    );
    let y = OrbitMass
        .with_position(position())
        .with_velocity(velocity())
        .with_mass(500.0);
    group.bench_function("two_body/OrbitMass", |b| {
        b.iter(|| state_derivative(black_box(0.0), black_box(y.vector), &assembly))
    });

    let assembly = Assembly::new(
        dynamics::build::<f64, 14, Full>(&two_body()),
        Context::new(epoch(), Frame::EARTH_J2000),
    );
    let y = Full
        .with_position(position())
        .with_velocity(velocity())
        .with_mass(500.0);
    group.bench_function("two_body/Full", |b| {
        b.iter(|| state_derivative(black_box(0.0), black_box(y.vector), &assembly))
    });

    group.finish();
}

/// What attaching a spacecraft costs, per derivative evaluation.
///
/// Nothing in this file's force list reads a mass or an area, so these differences
/// are overhead and nothing else. Two effects to separate: with a vehicle present
/// there is always a mass to divide a force by, so assembly takes the
/// `acceleration + force / mass` path even though the force is zero; and on a layout
/// carrying attitude, a vehicle supplies the inertia tensor that turns the
/// rotational half of the state from held-constant into Euler's equation.
fn bench_spacecraft(c: &mut Criterion) {
    let mut group = c.benchmark_group("derivative/spacecraft");

    let y = Orbit.with_position(position()).with_velocity(velocity());
    let bare = Assembly::new(
        dynamics::build::<f64, 6, Orbit>(&two_body()),
        Context::new(epoch(), Frame::EARTH_J2000),
    );
    group.bench_function("Orbit/without", |b| {
        b.iter(|| state_derivative(black_box(0.0), black_box(y.vector), &bare))
    });

    let equipped = Assembly::new(
        dynamics::build::<f64, 6, Orbit>(&two_body()),
        Context::new(epoch(), Frame::EARTH_J2000).with_spacecraft(spacecraft()),
    );
    group.bench_function("Orbit/with", |b| {
        b.iter(|| state_derivative(black_box(0.0), black_box(y.vector), &equipped))
    });

    let y = Full
        .with_position(position())
        .with_velocity(velocity())
        .with_mass(500.0)
        .with_angular_rate(angular_rate());
    let bare = Assembly::new(
        dynamics::build::<f64, 14, Full>(&two_body()),
        Context::new(epoch(), Frame::EARTH_J2000),
    );
    group.bench_function("Full/without", |b| {
        b.iter(|| state_derivative(black_box(0.0), black_box(y.vector), &bare))
    });

    let equipped = Assembly::new(
        dynamics::build::<f64, 14, Full>(&two_body()),
        Context::new(epoch(), Frame::EARTH_J2000).with_spacecraft(spacecraft()),
    );
    group.bench_function("Full/with", |b| {
        b.iter(|| state_derivative(black_box(0.0), black_box(y.vector), &equipped))
    });

    group.finish();

    // The same question at the propagation level, on one layout, so that the only
    // difference is the vehicle and not the dimension.
    let mut group = c.benchmark_group("propagation/spacecraft");
    group.sample_size(20);

    let forces = two_body();
    let period = Duration::from_seconds(kepler::period(SMA, MU));
    let spacecraft = spacecraft();

    group.bench_function("Orbit/without", |b| {
        b.iter(|| {
            let y0 = Orbit.with_position(position()).with_velocity(velocity());
            let mut propagation = Propagation::new(epoch(), y0, &forces)
                .duration(period)
                .accuracy(Accuracy::Medium);
            black_box(propagation.run().states.len())
        })
    });

    group.bench_function("Orbit/with", |b| {
        b.iter(|| {
            let y0 = Orbit.with_position(position()).with_velocity(velocity());
            let mut propagation = Propagation::new(epoch(), y0, &forces)
                .spacecraft(spacecraft.clone())
                .duration(period)
                .accuracy(Accuracy::Medium);
            black_box(propagation.run().states.len())
        })
    });

    group.finish();
}

/// One orbital period carrying attitude, against the same orbit without it.
///
/// The attitude adds eight components to the error control as well as the work of the
/// rotational derivative, so the step count can move too. It also brings the
/// shadow-set event, whose crossings cost a root search apiece, so this is the whole
/// price of carrying attitude rather than the derivative's share of it.
fn bench_attitude_propagation(c: &mut Criterion) {
    let mut group = c.benchmark_group("propagation/attitude");
    group.sample_size(20);

    let forces = two_body();
    let period = Duration::from_seconds(kepler::period(SMA, MU));
    let spacecraft = spacecraft();

    group.bench_function("orbit_only", |b| {
        b.iter(|| {
            let y0 = Orbit.with_position(position()).with_velocity(velocity());
            let mut propagation = Propagation::new(epoch(), y0, &forces)
                .duration(period)
                .accuracy(Accuracy::Medium);
            black_box(propagation.run().states.len())
        })
    });

    group.bench_function("with_attitude", |b| {
        b.iter(|| {
            let y0 = Full
                .with_position(position())
                .with_velocity(velocity())
                .with_mass(500.0)
                .with_angular_rate(angular_rate());
            let mut propagation = Propagation::new(epoch(), y0, &forces)
                .spacecraft(spacecraft.clone())
                .duration(period)
                .accuracy(Accuracy::Medium);
            black_box(propagation.run().states.len())
        })
    });

    group.finish();
}

/// A full propagation at each accuracy preset, over one day.
fn bench_propagation(c: &mut Criterion) {
    let mut group = c.benchmark_group("propagation/one_day");
    group.sample_size(20);

    let forces = two_body();
    for accuracy in [Accuracy::Low, Accuracy::Medium, Accuracy::High] {
        group.bench_function(format!("{accuracy:?}"), |b| {
            b.iter(|| {
                let y0 = Orbit.with_position(position()).with_velocity(velocity());
                let mut propagation = Propagation::new(epoch(), y0, &forces)
                    .duration(Duration::from_seconds(86400.0))
                    .accuracy(accuracy);
                black_box(propagation.run().states.len())
            })
        });
    }
    group.finish();
}

/// One orbital period, which is the unit most of the tests work in.
fn bench_one_period(c: &mut Criterion) {
    let forces = two_body();
    let period = Duration::from_seconds(kepler::period(SMA, MU));

    c.bench_function("propagation/one_period", |b| {
        b.iter(|| {
            let y0 = Orbit.with_position(position()).with_velocity(velocity());
            let mut propagation = Propagation::new(epoch(), y0, &forces)
                .duration(period)
                .accuracy(Accuracy::Medium);
            black_box(propagation.run().states.len())
        })
    });
}

/// The cost of a state transition matrix relative to a plain propagation.
///
/// Central differences would need thirteen propagations for a 6x6 and lose about
/// half the available digits. This tracks the ratio as the force models get
/// heavier.
#[cfg(feature = "autodiff")]
fn bench_state_transition_matrix(c: &mut Criterion) {
    use nalgebra::{Matrix6, Vector6};
    use num_dual::DualSVec64;

    let mut group = c.benchmark_group("propagation/stm");
    group.sample_size(20);

    let forces = two_body();
    let period = Duration::from_seconds(kepler::period(SMA, MU));
    let state = Vector6::new(
        position()[0],
        position()[1],
        position()[2],
        velocity()[0],
        velocity()[1],
        velocity()[2],
    );

    group.bench_function("plain", |b| {
        b.iter(|| {
            let y0 = Orbit.state(state);
            let mut propagation = Propagation::new(epoch(), y0, &forces)
                .duration(period)
                .accuracy(Accuracy::Medium);
            black_box(*propagation.run().states.last().unwrap())
        })
    });

    group.bench_function("with_stm", |b| {
        b.iter(|| {
            let (_, stm): (Vector6<f64>, Matrix6<f64>) = num_dual::jacobian(
                |y: Vector6<DualSVec64<6>>| {
                    let mut propagation = Propagation::new(epoch(), Orbit.state(y), &forces)
                        .duration(period)
                        .accuracy(Accuracy::Medium);
                    *propagation.run().states.last().unwrap()
                },
                &state,
            );
            black_box(stm)
        })
    });

    group.finish();
}

#[cfg(feature = "autodiff")]
criterion_group!(
    benches,
    bench_derivative,
    bench_spacecraft,
    bench_propagation,
    bench_one_period,
    bench_attitude_propagation,
    bench_state_transition_matrix
);

#[cfg(not(feature = "autodiff"))]
criterion_group!(
    benches,
    bench_derivative,
    bench_spacecraft,
    bench_propagation,
    bench_one_period,
    bench_attitude_propagation
);

criterion_main!(benches);