//! The finite burn.

use crate::error::{ManeuverError, ManeuverErrorKind};
use apsis_kepler::{GravityParameter, KeplerianElements, State};
use apsis_math::Vec3;

/// A maneuver that applies a constant acceleration for a duration.
///
/// A finite burn is the model for a thruster that is on for long enough that
/// its duration cannot be ignored: an electric thruster, whose burn may last
/// a large fraction of an orbit, or any burn where the spacecraft moves
/// appreciably while thrusting.
///
/// The acceleration is held constant in an inertial frame chosen by the
/// caller. That is the simplification this model makes: a real thruster
/// fires along a body axis that rotates as the spacecraft turns to hold a
/// direction, and modelling that requires the attitude history, which is not
/// part of this crate.
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct FiniteBurn {
    /// The acceleration vector, in kilometers per second squared.
    acceleration: Vec3,
    /// The duration, in seconds.
    duration: f64,
}

impl FiniteBurn {
    /// Creates a burn from an acceleration and a duration.
    ///
    /// Returns an error when a value is not finite, when the duration is not
    /// positive, or when the acceleration is zero.
    pub fn new(acceleration: Vec3, duration: f64) -> Result<Self, ManeuverError> {
        if !acceleration.is_finite() || !duration.is_finite() {
            return Err(ManeuverError::new(ManeuverErrorKind::NotFinite));
        }
        if duration <= 0.0 {
            return Err(ManeuverError::new(ManeuverErrorKind::NonPositiveDuration)
                .with_detail(format!("duration {duration}")));
        }
        if acceleration.norm() == 0.0 {
            return Err(
                ManeuverError::new(ManeuverErrorKind::NonPositiveAcceleration)
                    .with_detail("the acceleration is zero"),
            );
        }
        Ok(Self {
            acceleration,
            duration,
        })
    }

    /// Creates a burn of a given acceleration magnitude along a direction.
    ///
    /// The direction is normalized, so the acceleration magnitude is fixed
    /// whatever the length of the direction vector.
    pub fn along(direction: Vec3, acceleration: f64, duration: f64) -> Result<Self, ManeuverError> {
        if !acceleration.is_finite() || acceleration <= 0.0 {
            return Err(
                ManeuverError::new(ManeuverErrorKind::NonPositiveAcceleration)
                    .with_detail(format!("acceleration {acceleration}")),
            );
        }
        let unit = direction.normalize().ok_or_else(|| {
            ManeuverError::new(ManeuverErrorKind::InvalidState)
                .with_detail("the direction is zero or not finite")
        })?;
        Self::new(unit * acceleration, duration)
    }

    /// Returns the acceleration vector, in kilometers per second squared.
    pub const fn acceleration(self) -> Vec3 {
        self.acceleration
    }

    /// Returns the magnitude of the acceleration, in kilometers per second
    /// squared.
    pub fn acceleration_magnitude(self) -> f64 {
        self.acceleration.norm()
    }

    /// Returns the duration, in seconds.
    pub const fn duration(self) -> f64 {
        self.duration
    }

    /// Returns the velocity change the acceleration would give in free space,
    /// in kilometers per second.
    ///
    /// This is the acceleration times the duration, the delta-v the burn
    /// would deliver without gravity acting during it. Under gravity the
    /// achieved change is close to this for a short burn and differs from it
    /// for a long one, which is the difference the finite model exists to
    /// capture.
    pub fn nominal_delta_v(self) -> f64 {
        self.acceleration.norm() * self.duration
    }

    /// Returns the impulse equivalent to the nominal change.
    ///
    /// This is the impulsive burn the finite one approximates as its
    /// duration shrinks, and it is what [`FiniteBurn::apply`] returns in the
    /// limit of a constant gravity field.
    pub fn equivalent_impulse(self) -> Result<crate::impulse::Impulse, ManeuverError> {
        crate::impulse::Impulse::new(self.acceleration * self.duration)
    }

    /// Applies the burn to a state under two-body gravity.
    ///
    /// The position and velocity are integrated forward over the duration
    /// with the constant acceleration added to gravity. The exact solution
    /// of constant acceleration plus point mass gravity has no closed form,
    /// so the integration is a fixed step fourth order Runge-Kutta, which is
    /// accurate for the short arcs a burn occupies.
    ///
    /// The achieved velocity change is measured against a coast of the same
    /// duration, not against the initial velocity. That distinction matters:
    /// gravity changes the velocity during the burn even with no thrust, so
    /// the difference between the final and the initial velocity includes
    /// the natural rotation of the velocity vector. Comparing against the
    /// coast isolates the contribution of the thrust, which is what the
    /// burn bought; it also makes the result converge on the nominal change
    /// as the duration shrinks, which the difference from the initial
    /// velocity would not.
    ///
    /// Returns the state at the end of the burn together with the achieved
    /// velocity change.
    pub fn apply(
        self,
        state: State,
        parameter: GravityParameter,
        steps: usize,
    ) -> Result<BurnOutcome, ManeuverError> {
        if !state.position.is_finite() || !state.velocity.is_finite() {
            return Err(ManeuverError::new(ManeuverErrorKind::InvalidState));
        }
        if steps == 0 {
            return Err(ManeuverError::new(ManeuverErrorKind::InvalidState)
                .with_detail("the step count is zero"));
        }
        let mu = parameter.value();
        let step = self.duration / steps as f64;

        // The two-body acceleration at a position.
        let gravity = |position: Vec3| -> Vec3 {
            let radius = position.norm();
            if radius == 0.0 {
                return Vec3::ZERO;
            }
            position * (-mu / (radius * radius * radius))
        };
        // The thrust added to gravity, which is what the burn integrates.
        let thrust_and_gravity = |position: Vec3| -> Vec3 { gravity(position) + self.acceleration };

        // The burn trajectory.
        let mut position = state.position;
        let mut velocity = state.velocity;
        for _ in 0..steps {
            let (new_position, new_velocity) =
                rk4_step(position, velocity, step, &thrust_and_gravity);
            position = new_position;
            velocity = new_velocity;
        }

        // The coast trajectory over the same duration, with the same
        // integrator and step, so the only difference is the thrust.
        let mut coast_position = state.position;
        let mut coast_velocity = state.velocity;
        for _ in 0..steps {
            let (new_position, new_velocity) =
                rk4_step(coast_position, coast_velocity, step, &gravity);
            coast_position = new_position;
            coast_velocity = new_velocity;
        }

        if !position.is_finite() || !velocity.is_finite() {
            return Err(ManeuverError::new(ManeuverErrorKind::InvalidState)
                .with_detail("the integration diverged"));
        }

        Ok(BurnOutcome {
            state: State { position, velocity },
            achieved_delta_v: velocity - coast_velocity,
            nominal_delta_v: self.nominal_delta_v(),
            duration: self.duration,
        })
    }

    /// Applies the burn to a set of Keplerian elements.
    ///
    /// The elements are converted to a state, the burn is applied to it, and
    /// the result is converted back. Returns an error when the elements are
    /// not valid, when the integration fails, or when the resulting state is
    /// no longer an elliptical orbit, which happens when the burn is large
    /// enough to escape.
    pub fn apply_to_elements(
        self,
        elements: KeplerianElements,
        parameter: GravityParameter,
        steps: usize,
    ) -> Result<BurnOutcome, ManeuverError> {
        let state = elements.to_state(parameter).ok_or_else(|| {
            ManeuverError::new(ManeuverErrorKind::InvalidState)
                .with_detail("the elements are not a valid elliptical orbit")
        })?;
        self.apply(state, parameter, steps)
    }
}

/// What a finite burn produced.
///
/// The outcome keeps the final state, the change the burn achieved, and the
/// change the acceleration would have given in free space, so a caller can
/// see how much gravity cost or added during the burn.
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct BurnOutcome {
    /// The state at the end of the burn.
    pub state: State,
    /// The velocity change actually achieved, in kilometers per second.
    pub achieved_delta_v: Vec3,
    /// The velocity change the acceleration would give in free space, in
    /// kilometers per second.
    pub nominal_delta_v: f64,
    /// The burn duration, in seconds.
    pub duration: f64,
}

impl BurnOutcome {
    /// Returns the magnitude of the achieved velocity change.
    pub fn achieved_magnitude(&self) -> f64 {
        self.achieved_delta_v.norm()
    }

    /// Returns the difference between the achieved and nominal change, in
    /// kilometers per second.
    ///
    /// A negative value means gravity turned the velocity against the thrust
    /// during the burn, so the burn bought less than its free space delta-v.
    pub fn gravity_loss(&self) -> f64 {
        self.achieved_magnitude() - self.nominal_delta_v
    }

    /// Returns the elements at the end of the burn.
    ///
    /// Returns `None` when the final state is not an elliptical orbit, which
    /// happens when the burn is large enough to escape.
    pub fn elements(&self, parameter: GravityParameter) -> Option<KeplerianElements> {
        KeplerianElements::from_state(&self.state, parameter)
    }
}

/// One classical Runge-Kutta fourth order step of the two-body equations
/// with a constant acceleration added.
///
/// The state is `(position, velocity)` and the equations are the usual
/// `dr/dt = v` and `dv/dt = a(r)`, with `a` combining thrust and gravity.
/// The four stages are the standard ones; writing them out keeps the loop
/// allocation free.
fn rk4_step(
    position: Vec3,
    velocity: Vec3,
    step: f64,
    acceleration: &impl Fn(Vec3) -> Vec3,
) -> (Vec3, Vec3) {
    let k1_v = acceleration(position);
    let k1_r = velocity;

    let k2_v = acceleration(position + k1_r * (step / 2.0));
    let k2_r = velocity + k1_v * (step / 2.0);

    let k3_v = acceleration(position + k2_r * (step / 2.0));
    let k3_r = velocity + k2_v * (step / 2.0);

    let k4_v = acceleration(position + k3_r * step);
    let k4_r = velocity + k3_v * step;

    let new_position = position + (k1_r + k2_r * 2.0 + k3_r * 2.0 + k4_r) * (step / 6.0);
    let new_velocity = velocity + (k1_v + k2_v * 2.0 + k3_v * 2.0 + k4_v) * (step / 6.0);
    (new_position, new_velocity)
}