1 files changed, 31 insertions, 7 deletions
diff --git a/src/main.rs b/src/main.rs
index 92b0a1b..9f69b87 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,3 +1,4 @@
+use std::cmp::Ordering;
 use std::io::prelude::*;
 use std::io::stdin;
 
@@ -26,22 +27,45 @@ fn main() {
 }
 
 fn choose_command(state: &GameState) -> Command {
+    /* TODO:
+     * - Simulate both player and opponent moves
+     */
     let player_moves = state.valid_moves(0);
     let naive_result = player_moves
         .into_iter()
         .map(|player_move| {
             let mut state = state.clone();
-            state.update([player_move, Command::Nothing]);
+            state.update([player_move, Command::Accelerate]);
             (player_move, state)
         })
-        .max_by(|(_, state_1), (_, state_2)| {
-            state_1.players[0]
-                .position
-                .x
-                .cmp(&state_2.players[0].position.x)
-                .then(state_1.players[0].speed.cmp(&state_2.players[0].speed))
+        .flat_map(|(player_move, state)| {
+            state.valid_moves(0).into_iter().map(move |second_move| {
+                let mut second_move_state = state.clone();
+                second_move_state.update([second_move, Command::Accelerate]);
+                (player_move, second_move_state)
+            })
         })
+        .max_by(|(_, a), (_, b)| compare_states(a, b))
         .unwrap()
         .0;
     naive_result
 }
+
+fn compare_states(a: &GameState, b: &GameState) -> Ordering {
+    if a.status == GameStatus::PlayerOneWon && b.status == GameStatus::PlayerOneWon {
+        a.players[0].speed.cmp(&b.players[0].speed)
+    } else if a.status == GameStatus::PlayerOneWon {
+        Ordering::Greater
+    } else if b.status == GameStatus::PlayerOneWon {
+        Ordering::Less
+    } else {
+        let weighted_position_a = a.players[0].position.x + a.players[0].boosts * 2;
+        let weighted_position_b = b.players[0].position.x + b.players[0].boosts * 2;
+
+        weighted_position_a
+            .cmp(&weighted_position_b)
+            .then(a.players[0].speed.cmp(&b.players[0].speed))
+            .then(a.players[0].position.x.cmp(&b.players[0].position.x))
+            .then(a.players[0].boosts.cmp(&b.players[0].boosts))
+    }
+}