use std::collections::BTreeSet;
use std::time::Instant;
use bevy::app::App;
use bevy::color::palettes::basic::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()
})
.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 color = if thrust_solution.thrusters_on.contains(&thruster.1) {
LIMEGREEN
} else {
WHITE
};
gizmos.arrow_2d(
thruster.2.translation().xy(),
thruster.2.translation().xy() + thruster.2.rotation().mul_vec3(thruster.0.thrust_vector.extend(0.0)).xy(),
color
);
}
}
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>,
) {
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 */ }
debug!("input changed, recalculating thrust solution");
let start = Instant::now();
// 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;
if input.pressed(&ClientAction::ThrustForward) {
target_unit_vector += Vec3::new(0.0, -1.0, 0.0);
}
if input.pressed(&ClientAction::ThrustBackward) {
target_unit_vector += Vec3::new(0.0, 1.0, 0.0);
}
if input.pressed(&ClientAction::ThrustRight) {
target_unit_vector += Vec3::new(-1.0, 0.0, 0.0);
}
if input.pressed(&ClientAction::ThrustLeft) {
target_unit_vector += Vec3::new(1.0, 0.0, 0.0);
}
if input.pressed(&ClientAction::TorqueCw) {
target_unit_vector += Vec3::new(0.0, 0.0, -1.0);
}
if input.pressed(&ClientAction::TorqueCcw) {
target_unit_vector += Vec3::new(0.0, 0.0, 1.0);
}
if target_unit_vector == Vec3::ZERO {
debug!("no buttons are pressed; zeroing thrust solution");
solution.thrusters_on.clear();
debug!("solved thrust in {}ms", start.elapsed().as_millis());
return;
}
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;
};
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 {
warn!("issue while solving for thrust: part {:?} has no thrusters? skipping...", *part);
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.extend(0.0).cross(thruster_vector.extend(0.0)).z;
// magically assemble the vector!
let target_vector = thruster_vector.extend(thruster_torque);
all_thrusters.push((thruster_id, target_vector));
}
}
// calculate thrust and torque values
debug!("found {} thrusters, computing coefficients", all_thrusters.len());
let coefficients = all_thrusters.iter()
.map(|u| target_unit_vector.dot(u.1))
.collect::<Vec<_>>();
debug!("preparing model");
let mut problem = Problem::new(OptimizationDirection::Maximize);
// add variables to problem
let variables = coefficients.iter()
.map(|u| problem.add_binary_var(*u as f64))
.collect::<Vec<_>>();
debug!("prepared {} variables; solving", variables.len());
let ssolution = match 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;
}
}
}
};
debug!("found thrust solution!");
debug!("solution alignment (higher is better): {}", ssolution.objective());
let mut new_soln = ThrustSolution {
thrusters_on: BTreeSet::default()
};
for thruster in all_thrusters.iter().enumerate() {
debug!("solution: thruster #{} ({:?}): {}", thruster.0, thruster.1.0, ssolution.var_value_rounded(variables[thruster.0]));
if ssolution.var_value_rounded(variables[thruster.0]) == 1.0 {
new_soln.thrusters_on.insert(*thruster.1.0);
}
}
debug!("found thrust solution in {}ms", start.elapsed().as_millis());
*solution = new_soln;
return;
}