diff --git a/src/main/java/BattleShip/AIGridGUI.java b/src/main/java/BattleShip/AIGridGUI.java
index 57aaab0..042c5ba 100644
--- a/src/main/java/BattleShip/AIGridGUI.java
+++ b/src/main/java/BattleShip/AIGridGUI.java
@@ -333,6 +333,22 @@ public class AIGridGUI extends JPanel {
             if(killed) {
                 text = "Your " + s.getName() + " was sunk. Your turn.";
                 boolean unkilledCells = false;
+                for(BSButton bu : buttons) {
+                    //Mark killed cells.
+                    if(bu.getCellContents() == s) {
+                        bu.setBackground(Red);
+                        cellsKilled[bu.getGridLocation()] = true;
+                    }
+                    //Mark if any cell remains that has been hit but not yet killed.  If so, lock onto that cell.
+                    if(cellsHit[bu.getGridLocation()] && !cellsKilled[bu.getGridLocation()]) {
+                        f = bu.getGridLocation();
+                        unkilledCells = true;
+                    }
+                }
+                //If all hit cells have been killed, return to random guessing.
+                if(!unkilledCells) {
+                    randomGuess = true;
+                }
             }
         }
     }