this code works mostly fine, but when it cant find a 100% safe cell to open/click it junst clicks a cell where there 100% is a mine! instead of clicking the cell with lowest probability of mine being there. so your task is to rewrite the code, keep the input arguments same ofc so nothing breaks when i replace my code with yor new one. code:
class MinesweeperAI:
    def __init__(self, grid: List[List[str]]):
        """
        Initialize the Minesweeper AI with the current grid.
        """
        self.grid = grid
        self.height = len(grid)
        self.width = len(grid[0]) if self.height > 0 else 0
        self.safe_cells: Set[Tuple[int, int]] = set()
        self.mine_cells: Set[Tuple[int, int]] = set()
        self.cells_to_consider: Set[Tuple[int, int]] = set()
        self._init_knowledge()

    def _init_knowledge(self):
        """
        Initialize knowledge based on the current grid.

        """
        for r in range(self.height):
            for c in range(self.width):
                cell = self.grid[r][c]
                if cell.isdigit():
                    count = int(cell)
                    neighbors = self._get_neighbors(r, c)
                    covered = [n for n in neighbors if self.grid[n[0]][n[1]] == '□']
                    if count == 0:
                        # All neighboring covered cells are safe
                        for n in neighbors:
                            if self.grid[n[0]][n[1]] == '□':
                                self.safe_cells.add(n)
                    elif count == len(covered):
                        # All neighboring covered cells are mines
                        for n in covered:
                            self.mine_cells.add(n)
        self.cells_to_consider = self.safe_cells.copy()

    def _get_neighbors(self, row: int, col: int) -> Set[Tuple[int, int]]:
        """
        Get all valid neighboring cells for a given cell.
        """
        neighbors = set()
        for dr in [-1, 0, 1]:
            for dc in [-1, 0, 1]:
                if dr == 0 and dc == 0:
                    continue
                nr, nc = row + dr, col + dc
                if 0 <= nr < self.height and 0 <= nc < self.width:
                    neighbors.add((nr, nc))
        return neighbors

    def _update_knowledge(self):
        """
        Update the knowledge base by inferring new safe cells or mines.
        """
        changed = True
        while changed:
            changed = False
            for r in range(self.height):
                for c in range(self.width):
                    cell = self.grid[r][c]
                    if not cell.isdigit():
                        continue
                    count = int(cell)
                    neighbors = self._get_neighbors(r, c)
                    flagged = len([n for n in neighbors if n in self.mine_cells])
                    covered = [n for n in neighbors if self.grid[n[0]][n[1]] == '□' and n not in self.mine_cells]
                    if count == flagged:
                        # All covered cells are safe
                        for n in covered:
                            if n not in self.safe_cells:
                                self.safe_cells.add(n)
                                self.cells_to_consider.add(n)
                                changed = True
                    elif count == flagged + len(covered):
                        # All covered cells are mines
                        for n in covered:
                            if n not in self.mine_cells:
                                self.mine_cells.add(n)
                                changed = True

    def make_move(self) -> Tuple[int, int]:
        """
        Decide the next move to uncover.
        Returns the coordinates (row, col) of the cell to uncover.
        """
        self._update_knowledge()
        if self.cells_to_consider:
            # Choose a known safe cell
            return self.cells_to_consider.pop()

        # If no safe cells are known, use probability to guess
        probabilities = self._calculate_probabilities()
        if not probabilities:
            raise Exception("No moves left to make.")
        # Select the cell with the lowest probability of being a mine
        best_move = min(probabilities, key=lambda x: x[1])[0]
        return best_move

    def _calculate_probabilities(self) -> List[Tuple[Tuple[int, int], float]]:
        """
        Calculate the probability of each covered cell containing a mine.
        Returns a list of tuples containing cell coordinates and their mine probability.
        """
        probabilities = {}
        if stage == 1:
            total_mines = 10
        elif stage == 2:
            total_mines = 19
        elif stage == 3:
            total_mines = 32
        remaining_mines = total_mines - len(self.mine_cells)
        covered_cells = {(r, c) for r in range(self.height) for c in range(self.width)
                        if self.grid[r][c] == '□'}
        if not covered_cells:
            return []

        # Simple probability: remaining mines / remaining covered cells
        base_prob = remaining_mines / len(covered_cells) if len(covered_cells) > 0 else 1
        for cell in covered_cells:
            probabilities[cell] = base_prob

        # Refine probabilities based on the numbers
        for r in range(self.height):
            for c in range(self.width):
                cell = self.grid[r][c]
                if not cell.isdigit():
                    continue
                count = int(cell)
                neighbors = self._get_neighbors(r, c)
                flagged = len([n for n in neighbors if n in self.mine_cells])
                covered = [n for n in neighbors if self.grid[n[0]][n[1]] == '□' and n not in self.mine_cells]
                remaining = count - flagged
                if remaining <= 0:
                    continue
                prob = remaining / len(covered) if len(covered) > 0 else 1
                for n in covered:
                    if n in probabilities:
                        probabilities[n] = max(probabilities[n], prob)
        # Convert to list and sort
        prob_list = list(probabilities.items())
        prob_list.sort(key=lambda x: x[1])
        return prob_list

    def _estimate_total_mines(self) -> int:
        """
        Estimate the total number of mines in the grid.
        Adjust this method based on how total mines are defined.
        For this implementation, we'll assume a standard Minesweeper setting based on grid size.
        """
        # Example: Assume a 9x9 grid with 10 mines (standard Beginner)
        # Modify accordingly based on actual game settings
        # Here, we'll calculate mines based on the opened cells and their numbers
        total = 0
        for r in range(self.height):
            for c in range(self.width):
                cell = self.grid[r][c]
                if cell.isdigit():
                    total += int(cell)
        # Divide by 8 to estimate total mines (since each mine is counted up to 8 times)
        estimated = total // 2  # Heuristic
        return estimated if estimated > 0 else 10