Better move heuristics

2 files changed, 35 insertions, 8 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))
+    }
+}
diff --git a/src/state.rs b/src/state.rs
index ecbd408..99239ae 100644
--- a/src/state.rs
+++ b/src/state.rs
@@ -38,6 +38,10 @@ pub struct Position {
 
 impl GameState {
     pub fn update(&mut self, commands: [Command; 2]) {
+        if self.status != GameStatus::Continue {
+            return;
+        }
+
         self.do_command(0, &commands[0]);
         self.do_command(1, &commands[1]);
         self.update_player_collisions();
@@ -64,7 +68,6 @@ impl GameState {
         use Command::*;
         self.players[player_index].tick_boost();
 
-        // TODO: Command validation assertions
         match command {
             Nothing => {}
             Accelerate => self.players[player_index].accelerate(),