Implemented lattice restricted battleship searching

1 files changed, 49 insertions, 4 deletions

diff --git a/src/knowledge.rs b/src/knowledge.rs
index f009c10..5e45f07 100644
--- a/src/knowledge.rs
+++ b/src/knowledge.rs
@@ -51,10 +51,7 @@ impl Knowledge {
         let unknown_hits = self.opponent_map.cells_with_unknown_hits();
         let adjacent_cells = self.opponent_map.adjacent_unshot_cells(&unknown_hits);
       
-        let possible_placements = self.opponent_map
-            .ships
-            .values()
-            .flat_map(|knowledge| knowledge.possible_placements.clone());
+        let possible_placements = self.opponent_map.possible_placements();
 
         let mut max_score = 1;
         let mut best_cells = Vec::new();
@@ -74,6 +71,39 @@ impl Knowledge {
 
         best_cells
     }
+
+    pub fn get_most_possibility_shots_on_lattice(&self) -> Vec<Point> {
+        let on_lattice = self.opponent_map.cells_on_lattice(self.lattice_size());
+        let possible_placements = self.opponent_map.possible_placements();
+
+        let mut max_possibilities = 1;
+        let mut best_cells = Vec::new();
+
+        for cell in on_lattice {
+            let possibilities = possible_placements.iter()
+                .filter(|placement| placement.touches_point(cell)).count();
+            if possibilities > max_possibilities {
+                max_possibilities = possibilities;
+                best_cells = vec!(cell);
+            }
+            else if possibilities == max_possibilities {
+                best_cells.push(cell);
+            }
+        }
+
+        best_cells
+    }
+
+    fn lattice_size(&self) -> u16 {
+        let any_long_ships = self.opponent_map.ships.iter()
+            .any(|(ship, knowledge)| ship.length() >= 4 && !knowledge.destroyed);
+        if any_long_ships {
+            4
+        }
+        else {
+            2
+        }
+    }
 }
 
 #[derive(Serialize, Deserialize, Clone, Debug)]
@@ -173,6 +203,21 @@ impl OpponentMapKnowledge {
                 .filter(|&y| cells.iter().any(|z| z.is_adjacent(y)))
         }).collect()
     }
+
+    fn cells_on_lattice(&self, lattice_size: u16) -> Vec<Point> {
+        self.cells.iter().flat_map(|x| {
+            x.iter()
+                .filter(|y| !y.shot_attempted() && y.position.is_on_lattice(lattice_size))
+                .map(|y| y.position)
+        }).collect()
+    }
+
+    fn possible_placements(&self) -> Vec<PossibleShipPlacement> {
+        self.ships
+            .values()
+            .flat_map(|knowledge| knowledge.possible_placements.clone())
+            .collect()
+    }
 }
 
 #[derive(Serialize, Deserialize, Clone, Debug)]