use crate::rcs::Thruster;
use good_lp::{
    Expression, IntoAffineExpression, ProblemVariables, Solution, SolverModel, microlp, variable,
};
use log::error;
use nautilus_common::math::Vector;
use std::ops::Add;

#[derive(Copy, Clone, Debug)]
pub(super) struct BalancingConstants {
    /// Should be smaller than the other two constants
    throttle: f64,
    force: f64,
    torque: f64,
}

impl Default for BalancingConstants {
    fn default() -> Self {
        Self {
            throttle: 1e-3,
            force: 1.0,
            torque: 1.0,
        }
    }
}

#[allow(clippy::similar_names)]
pub(super) fn balance_thrust<const N: usize>(
    desired_force: Vector,
    desired_torque: Vector,
    thrusters: &[Thruster; N],
    constants: BalancingConstants,
) -> [f64; N] {
    let mut problem = ProblemVariables::new();
    let throttles: [_; N] = std::array::from_fn(|_| problem.add(variable().min(0).max(1)));

    let force_error_x_1 = problem.add(variable().min(0));
    let force_error_x_2 = problem.add(variable().min(0));
    let force_error_x = force_error_x_1 + force_error_x_2;
    let force_delta_x = force_error_x_1 - force_error_x_2;
    let force_error_y_1 = problem.add(variable().min(0));
    let force_error_y_2 = problem.add(variable().min(0));
    let force_error_y = force_error_y_1 + force_error_y_2;
    let force_delta_y = force_error_y_1 - force_error_y_2;
    let force_error_z_1 = problem.add(variable().min(0));
    let force_error_z_2 = problem.add(variable().min(0));
    let force_error_z = force_error_z_1 + force_error_z_2;
    let force_delta_z = force_error_z_1 - force_error_z_2;
    let force_error = force_error_x + force_error_y + force_error_z;

    let torque_error_x_1 = problem.add(variable().min(0));
    let torque_error_x_2 = problem.add(variable().min(0));
    let torque_error_x = torque_error_x_1 + torque_error_x_2;
    let torque_delta_x = torque_error_x_1 - torque_error_x_2;
    let torque_error_y_1 = problem.add(variable().min(0));
    let torque_error_y_2 = problem.add(variable().min(0));
    let torque_error_y = torque_error_y_1 + torque_error_y_2;
    let torque_delta_y = torque_error_y_1 - torque_error_y_2;
    let torque_error_z_1 = problem.add(variable().min(0));
    let torque_error_z_2 = problem.add(variable().min(0));
    let torque_error_z = torque_error_z_1 + torque_error_z_2;
    let torque_delta_z = torque_error_z_1 - torque_error_z_2;
    let torque_error = torque_error_x + torque_error_y + torque_error_z;

    let total_throttle = throttles.iter().fold(0.into_expression(), Expression::add);

    let generated_force_x = (0..N)
        .map(|i| thrusters[i].direction.x * thrusters[i].force * throttles[i])
        .fold(0.into_expression(), Expression::add);
    let generated_force_y = (0..N)
        .map(|i| thrusters[i].direction.y * thrusters[i].force * throttles[i])
        .fold(0.into_expression(), Expression::add);
    let generated_force_z = (0..N)
        .map(|i| thrusters[i].direction.z * thrusters[i].force * throttles[i])
        .fold(0.into_expression(), Expression::add);

    let torque_directions: [_; N] = std::array::from_fn(|i| {
        let cross = thrusters[i].position.cross_product(thrusters[i].direction);
        if cross == Vector::ZERO {
            Vector::ZERO
        } else {
            cross.normalize()
        }
    });
    let moment_arms: [f64; N] = std::array::from_fn(|i| thrusters[i].position.length());

    let generated_torque_x = (0..N)
        .map(|i| torque_directions[i].x * thrusters[i].force * throttles[i] * moment_arms[i])
        .fold(0.into_expression(), Expression::add);
    let generated_torque_y = (0..N)
        .map(|i| torque_directions[i].y * thrusters[i].force * throttles[i] * moment_arms[i])
        .fold(0.into_expression(), Expression::add);
    let generated_torque_z = (0..N)
        .map(|i| torque_directions[i].z * thrusters[i].force * throttles[i] * moment_arms[i])
        .fold(0.into_expression(), Expression::add);

    let to_minimize = (total_throttle.clone() * constants.throttle)
        + (force_error.clone() * constants.force)
        + (torque_error.clone() * constants.torque);

    let solution = problem
        .minimise(to_minimize)
        .using(microlp)
        .with(force_delta_x.eq(generated_force_x - desired_force.x))
        .with(force_delta_y.eq(generated_force_y - desired_force.y))
        .with(force_delta_z.eq(generated_force_z - desired_force.z))
        .with(torque_delta_x.eq(generated_torque_x - desired_torque.x))
        .with(torque_delta_y.eq(generated_torque_y - desired_torque.y))
        .with(torque_delta_z.eq(generated_torque_z - desired_torque.z))
        .solve();

    match solution {
        Ok(solution) => std::array::from_fn(|i| solution.value(throttles[i])),
        Err(err) => {
            error!("Failed to balance thrusters {err}");
            std::array::repeat(0.0)
        }
    }
}

#[cfg(test)]
mod tests {
    use crate::hardware::channelization::PinoutChannel;
    use crate::rcs::Thruster;
    use crate::rcs::thruster_balancing::{BalancingConstants, balance_thrust};
    use nautilus_common::math::Vector;

    const EPSILON: f64 = 1e-6;

    #[test]
    fn test_thruster_balancing() {
        // We can't fire the one thruster we have because it will induce a torque
        assert_eq!(
            [0.0],
            balance_thrust(
                Vector::X,
                Vector::ZERO,
                &[Thruster::new(
                    PinoutChannel::ExtA(0),
                    Vector::new(0.0, 1.0, 0.0),
                    Vector::X,
                    1.0
                ),],
                BalancingConstants::default(),
            )
        );

        // A single perfectly aligned thruster
        assert_eq!(
            [1.0],
            balance_thrust(
                Vector::X,
                Vector::ZERO,
                &[Thruster::new(
                    PinoutChannel::ExtA(0),
                    Vector::ZERO,
                    Vector::X,
                    1.0
                ),],
                BalancingConstants::default(),
            )
        );

        // Two balanced thrusters should both be fired
        assert_eq!(
            [1.0, 1.0],
            balance_thrust(
                Vector::X,
                Vector::ZERO,
                &[
                    Thruster::new(
                        PinoutChannel::ExtA(0),
                        Vector::new(0.0, 1.0, 0.0),
                        Vector::X,
                        0.5
                    ),
                    Thruster::new(
                        PinoutChannel::ExtA(0),
                        Vector::new(0.0, -1.0, 0.0),
                        Vector::X,
                        0.5
                    ),
                ],
                BalancingConstants::default(),
            )
        );

        // Two thrusters with different moment arms
        // Despite desiring a lot of force, we can't generate too much without inducing a torque
        assert!(
            balance_thrust(
                Vector::X * 10.0,
                Vector::ZERO,
                &[
                    Thruster::new(
                        PinoutChannel::ExtA(0),
                        Vector::new(0.0, 1.0, 0.0),
                        Vector::X,
                        0.5
                    ),
                    Thruster::new(
                        PinoutChannel::ExtA(0),
                        Vector::new(0.0, -2.0, 0.0),
                        Vector::X,
                        0.5
                    ),
                ],
                BalancingConstants::default(),
            )
            .into_iter()
            .zip([1.0, 0.5])
            .map(|(a, b)| a - b)
            .sum::<f64>()
            .abs()
                < EPSILON
        );

        // Four perfectly balanced thrusters
        assert!(
            balance_thrust(
                Vector::ZERO,
                Vector::Z * 2.0,
                &[
                    Thruster::new(
                        PinoutChannel::ExtA(0),
                        Vector::new(0.0, 1.0, 0.0),
                        Vector::X,
                        1.0
                    ),
                    Thruster::new(
                        PinoutChannel::ExtA(0),
                        Vector::new(0.0, 1.0, 0.0),
                        -Vector::X,
                        1.0
                    ),
                    Thruster::new(
                        PinoutChannel::ExtA(0),
                        Vector::new(0.0, -1.0, 0.0),
                        Vector::X,
                        1.0
                    ),
                    Thruster::new(
                        PinoutChannel::ExtA(0),
                        Vector::new(0.0, -1.0, 0.0),
                        -Vector::X,
                        1.0
                    ),
                ],
                BalancingConstants::default(),
            )
            .into_iter()
            .zip([0.0, 1.0, 1.0, 0.0])
            .map(|(a, b)| a - b)
            .sum::<f64>()
            .abs()
                < EPSILON
        );

        // Four perfectly balanced thrusters
        assert!(
            balance_thrust(
                Vector::X * 2.0,
                Vector::Z * 2.0,
                &[
                    Thruster::new(
                        PinoutChannel::ExtA(0),
                        Vector::new(0.0, 1.0, 0.0),
                        Vector::X,
                        1.0
                    ),
                    Thruster::new(
                        PinoutChannel::ExtA(0),
                        Vector::new(0.0, 1.0, 0.0),
                        -Vector::X,
                        1.0
                    ),
                    Thruster::new(
                        PinoutChannel::ExtA(0),
                        Vector::new(0.0, -1.0, 0.0),
                        Vector::X,
                        1.0
                    ),
                    Thruster::new(
                        PinoutChannel::ExtA(0),
                        Vector::new(0.0, -1.0, 0.0),
                        -Vector::X,
                        1.0
                    ),
                ],
                BalancingConstants::default(),
            )
            .into_iter()
            // Only the one thruster should fire - we can't do any better than that
            // any more thrusters we fire won't give us any more thrust
            .zip([0.0, 0.0, 1.0, 0.0])
            .map(|(a, b)| a - b)
            .sum::<f64>()
            .abs()
                < EPSILON
        );
    }
}
