use std::collections::BTreeSet;
use std::time::Instant;
use bevy::app::App;
use bevy::color::palettes::basic::{RED, WHITE};
use bevy::color::palettes::css::LIMEGREEN;
use bevy::math::Vec3Swizzles;
use leafwing_input_manager::prelude::ActionState;
use microlp::{OptimizationDirection, Problem};
use crate::attachment::Parts;
use crate::client::input::ClientAction;
use crate::ecs::thruster::{PartThrusters, Thruster, ThrusterOfPart};
use crate::prelude::*;
use crate::client::input::util::ActionStateExt;
use crate::ecs::Me;
use crate::thrust::ThrustSolution;
pub fn client_thrusters_plugin(app: &mut App) {
app
.insert_resource(ThrusterDebugRes(false))
.insert_resource(ThrustSolution {
thrusters_on: BTreeSet::default(),
converged: true,
})
.add_systems(Update, draw_thruster_debug)
.add_systems(Update, solve_thrust);
}
#[derive(Resource, Deref)]
pub struct ThrusterDebugRes(pub bool);
fn draw_thruster_debug(
thruster_debug_res: Res<ThrusterDebugRes>,
thrusters: Query<(&Thruster, Entity, &GlobalTransform)>,
thrust_solution: Res<ThrustSolution>,
mut gizmos: Gizmos,
) {
if !thruster_debug_res.0 { return };
for thruster in thrusters {
// Draw white if it's just a thruster, bright green if it's in the current thrust solution
let mut color = if thrust_solution.thrusters_on.contains(&thruster.1) {
LIMEGREEN
} else {
WHITE
};
// Exception: if the thrust solution failed to converge, RED
if !thrust_solution.converged {
color = RED;
}
let rescaled_thrust_vector = thruster.0.thrust_vector / 200.0;
gizmos.arrow_2d(
thruster.2.translation().xy(),
thruster.2.translation().xy() + thruster.2.rotation().mul_vec3(rescaled_thrust_vector.extend(0.0)).xy(),
color
);
}
}
// TODO: split this into two passes
fn solve_thrust(
me: Query<(Option<&Parts>, &GlobalTransform, Entity), With<Me>>,
parts: Query<&PartThrusters>,
thrusters: Query<(&Thruster, &GlobalTransform)>,
input: Res<ActionState<ClientAction>>,
mut solution: ResMut<ThrustSolution>,
mut events: MessageWriter<ThrustSolution>,
) {
if !(
input.button_changed(&ClientAction::ThrustForward)
|| input.button_changed(&ClientAction::ThrustBackward)
|| input.button_changed(&ClientAction::TorqueCw)
|| input.button_changed(&ClientAction::TorqueCcw)
|| input.button_changed(&ClientAction::ThrustRight)
|| input.button_changed(&ClientAction::ThrustLeft)
) { return; /* no changes, existing thrust solution is valid */ }
trace!("input changed, recalculating thrust solution");
let start = Instant::now();
solution.thrusters_on.clear();
solution.converged = false;
let Ok((our_parts, hearty_transform, hearty)) = me.single() else {
error!("could not solve for thrust: hearty does not exist?");
error!("failed to solve for thrust after {}ms", start.elapsed().as_millis());
return;
};
// determine our target vector:
// unit vector in the intended direction of movement
// Z-axis torque: this cursed thing is apparently standard
// +Z == counterclockwise/ccw
// -Z == clockwise/cw
let mut target_unit_vector = Vec3::ZERO;
let mut target_torque_vector = Vec3::ZERO;
let mut anything_pressed = false;
if input.pressed(&ClientAction::ThrustForward) {
anything_pressed = true;
target_unit_vector += hearty_transform.rotation() * Vec3::new(0.0, 1.0, 0.0);
}
if input.pressed(&ClientAction::ThrustBackward) {
anything_pressed = true;
target_unit_vector += hearty_transform.rotation() * Vec3::new(0.0, -1.0, 0.0);
}
if input.pressed(&ClientAction::ThrustRight) {
anything_pressed = true;
target_unit_vector += hearty_transform.rotation() * Vec3::new(1.0, 0.0, 0.0);
}
if input.pressed(&ClientAction::ThrustLeft) {
anything_pressed = true;
target_unit_vector += hearty_transform.rotation() * Vec3::new(-1.0, 0.0, 0.0);
}
if input.pressed(&ClientAction::TorqueCw) {
anything_pressed = true;
target_torque_vector += Vec3::new(0.0, 0.0, -1.0);
}
if input.pressed(&ClientAction::TorqueCcw) {
anything_pressed = true;
target_torque_vector += Vec3::new(0.0, 0.0, 1.0);
}
if !anything_pressed {
trace!("no buttons are pressed; zeroing thrust solution");
trace!("solved thrust in {}ms", start.elapsed().as_millis());
solution.converged = true;
events.write(solution.clone());
return;
}
if target_unit_vector != Vec3::ZERO {
target_unit_vector = target_unit_vector.normalize();
}
if target_torque_vector != Vec3::ZERO {
target_torque_vector = target_torque_vector.normalize();
}
let mut all_parts = vec![hearty];
if let Some(parts) = our_parts {
all_parts.extend(parts.iter());
}
// collect all thrusters on our ship, and figure out their thrust vectors
let mut all_thrusters = vec![];
for part in &all_parts {
let Ok(part_thrusters) = parts.get(*part) else {
continue;
};
for thruster_id in &**part_thrusters {
let Ok((thruster, thruster_transform)) = thrusters.get(*thruster_id) else {
warn!("issue while solving for thrust: thruster {:?} of part {:?} does not exist? skipping...", *thruster_id, *part);
continue;
};
// determine the thruster force in world space
let thruster_vector = thruster_transform.rotation().mul_vec3(thruster.thrust_vector.extend(0.0)).xy();
// determine our xy offset from hearty
let relative_translation = thruster_transform.translation().xy() - hearty_transform.translation().xy();
// determine our rotational offset from hearty
let relative_rotation = thruster_transform.rotation() * -hearty_transform.rotation();
let thruster_torque = relative_translation.normalize().extend(0.0).cross(thruster_vector.normalize().extend(0.0)).z;
let renormalized_thruster_torque = if thruster_torque.abs() < 0.1 {
0.0
} else if thruster_torque < 0.0 {
-1.0
} else {
1.0
};
trace!("thruster: {:?} {}({})", thruster_vector, thruster_torque, renormalized_thruster_torque);
// magically assemble the worldspace vector! for the solver (not shipspace)
all_thrusters.push((thruster_id,
thruster_vector.extend(0.0),
Vec3::new(0.0, 0.0, renormalized_thruster_torque)
));
}
}
// calculate thrust ~~and torque~~ values
trace!("found {} thrusters, computing coefficients", all_thrusters.len());
for thruster in &all_thrusters {
trace!("thruster on ship: {:?}", thruster);
}
let coefficients = all_thrusters.iter()
.map(|u| {
trace!("{} dot {}, {} dot {}", target_unit_vector, u.1.normalize(), target_torque_vector, u.2.normalize());
(
target_unit_vector.dot(u.1.normalize()),
target_torque_vector.dot(u.2.normalize())
)
})
.map(|u| {
trace!("=> {}, {}", u.0, u.1);
// improve reliability:
// if dot is <0.1, zap it entirely (this thruster is not helping)
// TODO: figure out how to make this adjustable
if u.0.abs() < 0.1 {
(0.0, u.1)
} else {
(u.0, u.1)
}
})
//.map(|u| u.powi(3))
.collect::<Vec<_>>();
trace!("preparing models");
let mut thrust_problem = Problem::new(OptimizationDirection::Maximize);
let mut torque_problem = Problem::new(OptimizationDirection::Maximize);
// add variables to problem
let variables = coefficients.iter()
.map(|u| {
(
thrust_problem.add_var(u.0 as f64, (0.0, 1.0)),
torque_problem.add_var(u.1 as f64, (0.0, 1.0))
)
})
.collect::<Vec<_>>();
trace!("prepared {} variables; solving", variables.len());
let thrust_solution = match thrust_problem.solve() {
Ok(soln) => soln,
Err(e) => {
match e {
microlp::Error::Infeasible => {
error!("failed to solve for thrust: constraints cannot be satisfied");
error!("failed to solve for thrust after {}ms", start.elapsed().as_millis());
return;
},
microlp::Error::Unbounded => {
error!("failed to solve for thrust: system is unbounded");
error!("failed to solve for thrust after {}ms", start.elapsed().as_millis());
return;
},
microlp::Error::InternalError(e) => {
error!("failed to solve for thrust: solver encountered internal error: {e}");
error!("failed to solve for thrust after {}ms", start.elapsed().as_millis());
return;
}
}
}
};
let torque_solution = match torque_problem.solve() {
Ok(soln) => soln,
Err(e) => {
match e {
microlp::Error::Infeasible => {
error!("failed to solve for torque: constraints cannot be satisfied");
error!("failed to solve for torque after {}ms", start.elapsed().as_millis());
return;
},
microlp::Error::Unbounded => {
error!("failed to solve for torque: system is unbounded");
error!("failed to solve for torque after {}ms", start.elapsed().as_millis());
return;
},
microlp::Error::InternalError(e) => {
error!("failed to solve for torque: solver encountered internal error: {e}");
error!("failed to solve for torque after {}ms", start.elapsed().as_millis());
return;
}
}
}
};
trace!("found thrust+torque solution!");
trace!("thrust solution alignment (higher is better): {}", thrust_solution.objective());
trace!("torque solution alignment (higher is better): {}", torque_solution.objective());
let mut new_soln = ThrustSolution {
thrusters_on: BTreeSet::default(),
converged: true
};
for thruster in all_thrusters.iter().enumerate() {
trace!("thrust solution: thruster #{} ({:?}): {} @ coeff {}", thruster.0, thruster.1.0, thrust_solution.var_value(variables[thruster.0].0), coefficients[thruster.0].0);
trace!("torque solution: thruster #{} ({:?}): {} @ coeff {}", thruster.0, thruster.1.0, torque_solution.var_value(variables[thruster.0].1), coefficients[thruster.0].1);
// TODO: make this more easily adjustable
if *thrust_solution.var_value(variables[thruster.0].0) > 0.8 || *torque_solution.var_value(variables[thruster.0].1) > 0.8 {
new_soln.thrusters_on.insert(*thruster.1.0);
}
}
let elapsed = start.elapsed();
debug!(?elapsed, ?target_unit_vector, ?target_torque_vector, "solved for thrust and torque");
*solution = new_soln;
events.write(solution.clone());
return;
}