Allowed monte carlo search to use iron curtains

2 files changed, 14 insertions, 4 deletions

diff --git a/src/engine/bitwise_engine.rs b/src/engine/bitwise_engine.rs
index e9c1b0a..fb46567 100644
--- a/src/engine/bitwise_engine.rs
+++ b/src/engine/bitwise_engine.rs
@@ -87,6 +87,13 @@ impl BitwiseGameState {
     pub fn player_has_max_teslas(&self) -> bool { self.player_buildings.count_teslas() >= TESLA_MAX }
     pub fn opponent_has_max_teslas(&self) -> bool { self.opponent_buildings.count_teslas() >= TESLA_MAX }
 
+    pub fn player_can_build_iron_curtain(&self) -> bool {
+        self.player_buildings.iron_curtain_available && self.player_buildings.iron_curtain_remaining == 0 && self.player.energy >= IRON_CURTAIN_PRICE
+    }
+    pub fn opponent_can_build_iron_curtain(&self) -> bool {
+        self.opponent_buildings.iron_curtain_available && self.opponent_buildings.iron_curtain_remaining == 0 && self.opponent.energy >= IRON_CURTAIN_PRICE
+    }
+
     pub fn unoccupied_player_cell_count(&self) -> usize { self.player_buildings.occupied.count_zeros() as usize }
     pub fn unoccupied_opponent_cell_count(&self) -> usize { self.opponent_buildings.occupied.count_zeros() as usize }
     pub fn location_of_unoccupied_player_cell(&self, i: usize) -> Point  {
diff --git a/src/strategy/monte_carlo.rs b/src/strategy/monte_carlo.rs
index 57eaaff..95d87fa 100644
--- a/src/strategy/monte_carlo.rs
+++ b/src/strategy/monte_carlo.rs
@@ -116,21 +116,24 @@ fn simulate_to_endstate<R: Rng>(command_score: &mut CommandScore, state: &Bitwis
 
 fn random_player_move<R: Rng>(state: &BitwiseGameState, rng: &mut R) -> Command {
     let all_buildings = sensible_buildings(&state.player, &state.player_buildings, state.player_has_max_teslas());
-    random_move(&all_buildings, rng, state.unoccupied_player_cell_count(), |i| state.location_of_unoccupied_player_cell(i))
+    random_move(&all_buildings, state.player_can_build_iron_curtain(), rng, state.unoccupied_player_cell_count(), |i| state.location_of_unoccupied_player_cell(i))
 }
 
 fn random_opponent_move<R: Rng>(state: &BitwiseGameState, rng: &mut R) -> Command {
     let all_buildings = sensible_buildings(&state.opponent, &state.opponent_buildings, state.opponent_has_max_teslas());
-    random_move(&all_buildings, rng, state.unoccupied_opponent_cell_count(), |i| state.location_of_unoccupied_opponent_cell(i))
+    random_move(&all_buildings, state.opponent_can_build_iron_curtain(), rng, state.unoccupied_opponent_cell_count(), |i| state.location_of_unoccupied_opponent_cell(i))
 }
 
 // TODO: Given enough energy, most opponents won't do nothing
-fn random_move<R: Rng, F:Fn(usize)->Point>(all_buildings: &[BuildingType], rng: &mut R, free_positions_count: usize, get_point: F) -> Command {
+fn random_move<R: Rng, F:Fn(usize)->Point>(all_buildings: &[BuildingType], iron_curtain_available: bool, rng: &mut R, free_positions_count: usize, get_point: F) -> Command {
     let nothing_count = 1;
-    let building_choice_index = rng.gen_range(0, nothing_count + all_buildings.len());
+    let iron_curtain_count = if iron_curtain_available { 1 } else { 0 };
+    let building_choice_index = rng.gen_range(0, all_buildings.len() + nothing_count + iron_curtain_count);
     
     if building_choice_index == all_buildings.len() {
         Command::Nothing
+    } else if iron_curtain_available && building_choice_index == all_buildings.len() + 1 {
+        Command::IronCurtain
     } else {
         let position_choice = rng.gen_range(0, free_positions_count);
         Command::Build(