Event-based heuristic at the beginning, astar at the end

4 files changed, 130 insertions, 51 deletions

diff --git a/src/lib.rs b/src/lib.rs
index 04629c3..519b804 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -10,31 +10,57 @@ use pathfinding::prelude::*;
 use state::*;
 use std::cmp::Ordering;
 
-pub fn choose_command(state: &GameState) -> Command {
-    // choose_command_with_looking_forward_heuristic(state)
-    choose_command_with_astar(state)
+pub fn choose_command(round: usize, state: &GameState) -> Command {
+    if round <= 100 {
+        choose_command_with_looking_forward_heuristic(state)
+    } else {
+        choose_command_with_astar(state)
+    }
+}
+
+#[derive(Debug, Clone)]
+struct HeuristicLookaheadState {
+    moves: Vec<Command>,
+    events: GameStateUpdateEvents,
+    current_state: GameState,
 }
 
 fn choose_command_with_looking_forward_heuristic(state: &GameState) -> Command {
-    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::Accelerate]);
-            (player_move, state)
-        })
-        .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)
+    let depth = 3;
+    let mut states = vec![HeuristicLookaheadState {
+        moves: Vec::new(),
+        events: GameStateUpdateEvents::default(),
+        current_state: state.clone(),
+    }];
+
+    for _ in 0..depth {
+        states = states
+            .into_iter()
+            .flat_map(move |state| {
+                state
+                    .current_state
+                    .good_moves(0)
+                    .into_iter()
+                    .map(move |player_move| {
+                        let mut state = state.clone();
+                        state.moves.push(player_move);
+                        state
+                            .current_state
+                            .update([player_move, Command::Accelerate], &mut state.events);
+                        state
+                    })
             })
-        })
-        .max_by(|(_, a), (_, b)| compare_states(a, b))
+            .collect();
+    }
+
+    states
+        .into_iter()
+        .max_by(|a, b| compare_events(&a.events, &b.events))
+        .unwrap()
+        .moves
+        .into_iter()
+        .next()
         .unwrap()
-        .0;
-    naive_result
 }
 
 fn compare_states(a: &GameState, b: &GameState) -> Ordering {
@@ -56,7 +82,20 @@ fn compare_states(a: &GameState, b: &GameState) -> Ordering {
     }
 }
 
+fn compare_events(a: &GameStateUpdateEvents, b: &GameStateUpdateEvents) -> Ordering {
+    let a = &a.players[0];
+    let b = &b.players[0];
+    a.boosts_collected
+        .cmp(&b.boosts_collected)
+        .then(a.boosts_maintained.cmp(&b.boosts_maintained))
+        .then(a.distance_travelled.cmp(&b.distance_travelled))
+        .then(a.mud_hit.cmp(&b.mud_hit).reverse())
+        .then(a.boosts_used.cmp(&b.boosts_used).reverse())
+}
+
 fn choose_command_with_astar(state: &GameState) -> Command {
+    // TODO: Find all shortest paths, choose the one that has the
+    // highest end speed, or otherwise wins
     shortest_path_first_command(state).unwrap_or(Command::Accelerate)
 }
 
@@ -70,7 +109,10 @@ fn shortest_path_first_command(initial_state: &GameState) -> Option<Command> {
                 .filter(|player_move| *player_move != Command::UseOil)
                 .map(|player_move| {
                     let mut state = state.clone();
-                    state.update([player_move, Command::Decelerate]);
+                    state.update(
+                        [player_move, Command::Decelerate],
+                        &mut GameStateUpdateEvents::default(),
+                    );
                     (state, 1)
                 })
                 .collect::<Vec<_>>()
@@ -91,7 +133,10 @@ fn shortest_path_first_command(initial_state: &GameState) -> Option<Command> {
                 .filter(|player_move| *player_move != Command::UseOil)
                 .find(|player_move| {
                     let mut state = state.clone();
-                    state.update([*player_move, Command::Decelerate]);
+                    state.update(
+                        [*player_move, Command::Decelerate],
+                        &mut GameStateUpdateEvents::default(),
+                    );
                     state == *next
                 })
                 .unwrap();
diff --git a/src/main.rs b/src/main.rs
index a652b8e..c5f7857 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -6,11 +6,14 @@ use vroomba::*;
 
 fn main() {
     for line in stdin().lock().lines() {
-        let round_number = line.expect("Failed to read line from stdin: {}");
+        let round_number = line
+            .expect("Failed to read line from stdin: {}")
+            .parse::<usize>()
+            .expect("Round number was not an unsigned integer: {}");
         let command =
             match json::read_state_from_json_file(&format!("./rounds/{}/state.json", round_number))
             {
-                Ok(state) => choose_command(&state),
+                Ok(state) => choose_command(round_number, &state),
                 Err(e) => {
                     eprintln!("WARN: State file could not be parsed: {}", e);
                     Command::Nothing
diff --git a/src/state.rs b/src/state.rs
index 7cc1ab4..be42228 100644
--- a/src/state.rs
+++ b/src/state.rs
@@ -12,6 +12,8 @@ pub enum GameStatus {
     Draw, // Until I add score I guess
 }
 
+// TODO: Maintain sorted vecs for these BTrees? Do the range counts
+// with binary searches to find indices only.
 #[derive(Debug, Clone, Hash, PartialEq, Eq)]
 pub struct GameState {
     pub status: GameStatus,
@@ -43,27 +45,26 @@ pub struct GameStateUpdateEvents {
 }
 #[derive(Debug, Clone, Default)]
 pub struct GameStatePlayerUpdateEvents {
-    pub boosts_collected: u16,
     pub boosts_used: u16,
+    pub boosts_collected: u16,
     pub boosts_maintained: u8,
     pub mud_hit: u16,
     pub oil_collected: u16,
     pub distance_travelled: u16,
-    pub rounds_to_finish: Option<u16>,
 }
 
 impl GameState {
-    pub fn update(&mut self, commands: [Command; 2]) -> GameStateUpdateEvents {
+    pub fn update(&mut self, commands: [Command; 2], events: &mut GameStateUpdateEvents) {
         if self.status != GameStatus::Continue {
-            return GameStateUpdateEvents::default();
+            return;
         }
 
         let next_positions = [
-            self.do_command(0, &commands[0]),
-            self.do_command(1, &commands[1]),
+            self.do_command(0, &commands[0], &mut events.players[0]),
+            self.do_command(1, &commands[1], &mut events.players[1]),
         ];
-        let next_positions = self.update_player_collisions(next_positions);
-        self.update_player_travel(next_positions);
+        let next_positions = self.update_player_collisions(next_positions, events);
+        self.update_player_travel(next_positions, events);
 
         self.status = if self.players[0].finished() && self.players[1].finished() {
             if self.players[0].speed > self.players[1].speed {
@@ -80,9 +81,6 @@ impl GameState {
         } else {
             GameStatus::Continue
         };
-
-        // TODO
-        GameStateUpdateEvents::default()
     }
 
     pub fn reset_players_to_start(&mut self) {
@@ -96,7 +94,12 @@ impl GameState {
         }
     }
 
-    fn do_command(&mut self, player_index: usize, command: &Command) -> Position {
+    fn do_command(
+        &mut self,
+        player_index: usize,
+        command: &Command,
+        events: &mut GameStatePlayerUpdateEvents,
+    ) -> Position {
         use Command::*;
         self.players[player_index].tick_boost();
         let mut next_y = self.players[player_index].position.y;
@@ -107,7 +110,10 @@ impl GameState {
             Decelerate => self.players[player_index].decelerate(),
             TurnLeft => next_y = next_y.saturating_sub(1).max(MIN_Y),
             TurnRight => next_y = next_y.saturating_add(1).min(MAX_Y),
-            UseBoost => self.players[player_index].boost(),
+            UseBoost => {
+                events.boosts_used += 1;
+                self.players[player_index].boost();
+            }
             UseOil => {
                 debug_assert!(self.players[player_index].oils > 0);
                 self.players[player_index].oils = self.players[player_index].oils.saturating_sub(1);
@@ -133,7 +139,11 @@ impl GameState {
         }
     }
 
-    fn update_player_collisions(&mut self, next_positions: [Position; 2]) -> [Position; 2] {
+    fn update_player_collisions(
+        &mut self,
+        next_positions: [Position; 2],
+        _events: &mut GameStateUpdateEvents,
+    ) -> [Position; 2] {
         let same_lanes_before = self.players[0].position.y == self.players[1].position.y;
         let same_lanes_after = next_positions[0].y == next_positions[1].y;
         let first_passing_second = self.players[0].position.x < self.players[1].position.x
@@ -174,14 +184,24 @@ impl GameState {
         }
     }
 
-    fn update_player_travel(&mut self, next_positions: [Position; 2]) {
-        for (player, next_position) in self.players.iter_mut().zip(next_positions.iter()) {
+    fn update_player_travel(
+        &mut self,
+        next_positions: [Position; 2],
+        events: &mut GameStateUpdateEvents,
+    ) {
+        for (i, (player, next_position)) in self
+            .players
+            .iter_mut()
+            .zip(next_positions.iter())
+            .enumerate()
+        {
             player.move_along(
                 *next_position,
                 &self.muds,
                 &self.oil_spills,
                 &self.powerup_oils,
                 &self.powerup_boosts,
+                &mut events.players[i],
             );
         }
     }
@@ -290,6 +310,7 @@ impl Player {
         oil_spills: &BTreeSet<Position>,
         powerup_oils: &BTreeSet<Position>,
         powerup_boosts: &BTreeSet<Position>,
+        events: &mut GameStatePlayerUpdateEvents,
     ) {
         let range = (
             Included(Position {
@@ -302,18 +323,28 @@ impl Player {
             }),
         );
 
-        for _ in muds.range(range) {
+        let mud_hit = muds
+            .range(range)
+            .count()
+            .saturating_add(oil_spills.range(range).count());
+        for _ in 0..mud_hit {
             self.decelerate_from_obstacle();
         }
-        for _ in oil_spills.range(range) {
-            self.decelerate_from_obstacle();
+
+        let oil_collected = powerup_oils.range(range).count() as u16;
+        self.oils = self.oils.saturating_add(oil_collected);
+        let boosts_collected = powerup_boosts.range(range).count() as u16;
+        self.boosts = self.boosts.saturating_add(boosts_collected);
+
+        events.mud_hit = events.mud_hit.saturating_add(mud_hit as u16);
+        events.boosts_collected = events.boosts_collected.saturating_add(boosts_collected);
+        events.oil_collected = events.oil_collected.saturating_add(oil_collected);
+        events.distance_travelled = events
+            .distance_travelled
+            .saturating_add(next_position.x.saturating_sub(self.position.x));
+        if self.speed == SPEED_BOOST {
+            events.boosts_maintained = events.boosts_maintained.saturating_add(1);
         }
-        self.oils = self
-            .oils
-            .saturating_add(powerup_oils.range(range).count() as u16);
-        self.boosts = self
-            .boosts
-            .saturating_add(powerup_boosts.range(range).count() as u16);
 
         self.position = next_position;
     }
diff --git a/tests/live_comparison.rs b/tests/live_comparison.rs
index 54ac1f7..3e29069 100644
--- a/tests/live_comparison.rs
+++ b/tests/live_comparison.rs
@@ -2,7 +2,7 @@ extern crate vroomba;
 
 use vroomba::command::Command;
 use vroomba::json;
-use vroomba::state::GameState;
+use vroomba::state::{GameState, GameStateUpdateEvents};
 
 use std::fs::File;
 use std::io::prelude::*;
@@ -48,7 +48,7 @@ fn test_from_replay(replay_folder: &Path) {
         state.players[1].oils = 1;
 
         println!("Player 1: {}, Player 2: {}", player, opponent);
-        state.update([player, opponent]);
+        state.update([player, opponent], &mut GameStateUpdateEvents::default());
 
         println!("State {}: {:?}", i + 1, state);
         assert_eq!(