mine_sweeper_solver.py

from collections import defaultdict
from collections.abc import Iterable
from dataclasses import dataclass
from enum import Enum
from itertools import chain
import random


class ActionType(Enum):
    """Enum for different action types"""

    CLICK = "click"
    FLAG = "flag"
    UNFLAG = "unflag"
    CHORD = "chord"
    HINT = "hint"


@dataclass
class Cell:
    """Dataclass representing cell information"""

    row: int
    col: int
    value: int  # -1 if not revealed/flagged, else adjacentMines count
    isRevealed: bool  # camelCase to match game data
    isFlagged: bool  # camelCase to match game data
    hint: str | None = None
    hintDetail: str | None = None


@dataclass
class Action:
    """Dataclass representing an action to take"""

    type: ActionType  # camelCase to match game conventions
    row: int
    col: int
    hint: str | None = None  # camelCase
    hintDetail: str | None = None  # camelCase


def getNeighbors(cell: Cell, cells: list[list[Cell]]) -> list[Cell]:
    """Get all valid neighbors of a cell"""
    neighbors = []
    for di in [-1, 0, 1]:
        for dj in [-1, 0, 1]:
            if di == 0 and dj == 0:
                continue
            ni, nj = cell.row + di, cell.col + dj
            if 0 <= ni < len(cells) and 0 <= nj < len(cells[0]):
                neighbors.append(cells[ni][nj])
    return neighbors


def filterUnrevealedNeighbors(neighbors: list[Cell]) -> list[Cell]:
    return [n for n in neighbors if not n.isRevealed and not n.isFlagged]


def countFlagged(neighbors: Iterable[Cell]) -> int:
    return sum(1 for n in neighbors if n.isFlagged)


def isRevealed(cell: Cell) -> bool:
    return cell.isRevealed and cell.value >= 0 and not cell.isFlagged


def isUnreached(cell: Cell, neighbors: Iterable[Cell]) -> bool:
    if cell.isRevealed or cell.isFlagged:
        return False
    return all(not n.isRevealed for n in neighbors)


def filterUnreachedCells(cells: list[list[Cell]]) -> list[Cell]:
    return [
        cell
        for cell in chain.from_iterable(cells)
        if isUnreached(cell, getNeighbors(cell, cells))
    ]


def flagIfMatch(cell: Cell, neighbors: list[Cell]) -> list[Action]:
    """Flag neighbors if they match the cell's value."""
    if not isRevealed(cell):
        return []

    unrevealedCells = filterUnrevealedNeighbors(neighbors)
    if not unrevealedCells:
        return []

    flaggedCount = countFlagged(neighbors)
    if len(unrevealedCells) != (cell.value - flaggedCount):
        return []
    return [Action(ActionType.FLAG, n.row, n.col) for n in unrevealedCells]


def revealIfSafe(cell: Cell, neighbors: list[Cell]) -> list[Action]:
    """Reveal neighbors if they are safe."""
    if not isRevealed(cell):
        return []

    unrevealedCells = filterUnrevealedNeighbors(neighbors)
    if not unrevealedCells:
        return []

    flaggedCount = countFlagged(neighbors)
    if cell.value != flaggedCount:
        return []
    return [Action(ActionType.CLICK, n.row, n.col) for n in unrevealedCells]


# ---
# Advanced solving methods


@dataclass
class CellExtended:
    row: int
    col: int
    mineCount: int  # -1 for unrevealed, this is the value that already minus flagged count

    neighbors: list["CellExtended"]  # neighbors is speical here
    # for unrevealed cell -> neighbors of working cells
    # for working cell -> neighbors of unrevealed cells
    # The cellExtended will only has this two type of cell

    def __hash__(self):
        return hash((self.row, self.col))

    def toCellStr(self, other_str: str = "") -> str:
        return f"Cell({self.row}, {self.col}, {self.mineCount}{other_str})"

    def __str__(self) -> str:
        neighborStr = ",".join(map(lambda c: c.toCellStr(), self.neighbors))
        return self.toCellStr(f", neighbors=[{neighborStr}]")

    def __repr__(self) -> str:
        return self.__str__()


def toCellExtended(cells: list[list[Cell]]) -> list[CellExtended]:
    # extract to new Cell type
    newCellsMap: dict[tuple[int, int], CellExtended] = {}

    def getOrCreateCellExtended(cell: Cell) -> CellExtended:
        key = (cell.row, cell.col)
        if key not in newCellsMap:
            newCellsMap[key] = CellExtended(cell.row, cell.col, -1, [])
        return newCellsMap[key]

    for row in cells:
        for cell in row:
            # Find revealed but unsolved cells
            if not isRevealed(cell):
                continue
            neighbors = getNeighbors(cell, cells)
            unrevealedNeighbors = filterUnrevealedNeighbors(neighbors)
            if len(unrevealedNeighbors) == 0:
                continue

            # ensure the board can not simply be solved
            remainMine = cell.value - countFlagged(neighbors)
            assert remainMine > 0
            newCell = getOrCreateCellExtended(cell)
            newCell.mineCount = remainMine

            for unrevealCell in unrevealedNeighbors:
                unreveal = getOrCreateCellExtended(unrevealCell)
                assert unreveal.mineCount == -1  # double check
                unreveal.neighbors.append(newCell)
                newCell.neighbors.append(unreveal)

    newCells: list[CellExtended] = []
    for cell in newCellsMap.values():
        cell.neighbors = list(set(cell.neighbors))
        newCells.append(cell)
    return newCells


def filterWorkingCellExtended(cells: list[CellExtended]) -> list[CellExtended]:
    return [cell for cell in cells if cell.mineCount >= 0]


def filterUnrevealedCellExtended(cells: list[CellExtended]) -> list[CellExtended]:
    return [cell for cell in cells if cell.mineCount == -1]


# utils
def iterPermutation(n: int, k: int) -> list[list[bool]]:
    """N=number of items, K=number of false"""
    if n < k or k < 0:
        assert False, "Invalid parameters n={}, k={}".format(n, k)  # debug
        # return []

    result = []

    def recursivePermutation(n: int, k: int, current_results: list[bool]):
        if k == 0:
            result.append([*current_results, *([False] * n)])
            return
        if k == n:
            result.append([*current_results, *([True] * n)])
            return
        recursivePermutation(n - 1, k - 1, current_results + [True])
        recursivePermutation(n - 1, k, current_results + [False])

    recursivePermutation(n, k, [])
    return result


@dataclass
class PermutationCount:
    mineCount: int
    safeCount: int
    part: int

    def mineProb(self) -> float:
        return self.mineCount / (self.mineCount + self.safeCount)

    def add(self, isSafe: bool, count: int):
        if isSafe:
            self.safeCount += count
        else:
            self.mineCount += count


def calcPermutationCounts(
    guess: dict[CellExtended, bool],
    unhandledWorkingCells: set[CellExtended],
    handledWorkingCells: set[CellExtended],
    permutationCounts: dict[CellExtended, PermutationCount],
    remainingMines: int,
    checkAnyValid: bool = False,
) -> int:
    #  guess: true = safe, false = mine, None = unguessed
    # return number of solutions

    # success: all handled
    if remainingMines < 0:
        return 0
    if not unhandledWorkingCells:
        return 1

    cell = next(iter(unhandledWorkingCells))
    # print("Handling cell:", cell)

    # check current guess is valid based on this cell
    unguessCells: list[CellExtended] = []
    guessFlagCount = 0
    for unrevealedCell in cell.neighbors:
        if unrevealedCell in guess:
            guessFlagCount += guess[unrevealedCell] == False
        else:
            unguessCells.append(unrevealedCell)
    unguessCount = len(unguessCells)
    if guessFlagCount > cell.mineCount:
        return 0
    if guessFlagCount + unguessCount < cell.mineCount:
        return 0

    usedMines = cell.mineCount - guessFlagCount
    if usedMines > remainingMines:
        return 0

    # find related working cells from unguessCells
    relatedWorkingCells: set[CellExtended] = set()
    for unguessCell in unguessCells:
        for relatedCell in unguessCell.neighbors:
            if (
                relatedCell not in handledWorkingCells
                and relatedCell not in unhandledWorkingCells
            ):
                relatedWorkingCells.add(relatedCell)

    # update unhandled and handled
    unhandledWorkingCells.remove(cell)
    handledWorkingCells.add(cell)
    for relatedCell in relatedWorkingCells:
        unhandledWorkingCells.add(relatedCell)

    resultCount = 0
    hasTryGuess = False
    for i, permutation in enumerate(
        iterPermutation(unguessCount, unguessCount - usedMines)
    ):
        hasTryGuess = True
        for unguessCell, isSafe in zip(unguessCells, permutation):
            guess[unguessCell] = isSafe
            # print(f"[{cell}] Try {i} Guess {unguessCell} as {'safe' if isSafe else 'mine'}")

        result = calcPermutationCounts(
            guess,
            unhandledWorkingCells,
            handledWorkingCells,
            permutationCounts,
            remainingMines - usedMines,
            checkAnyValid=checkAnyValid,
        )
        for unguessCell in unguessCells:
            permutationCounts[unguessCell].add(guess[unguessCell], result)

        resultCount += result

        # early stop, if more than one valid solution found
        if checkAnyValid and result > 0:
            break

    # remove guess
    if hasTryGuess:
        for unguessCell in unguessCells:
            del guess[unguessCell]

    # update unhandled and handled back
    unhandledWorkingCells.add(cell)
    handledWorkingCells.remove(cell)
    for relatedCell in relatedWorkingCells:
        unhandledWorkingCells.remove(relatedCell)
    return resultCount


def calcActionByFalseGuessing(cells: list[list[Cell]], totalMines: int) -> list[Action]:
    totalFlags = countFlagged(chain.from_iterable(cells))
    newCells = toCellExtended(cells)
    permutationMap = defaultdict(lambda: PermutationCount(0, 0, 0))

    actions = []
    for unrevealedCell in filterUnrevealedCellExtended(newCells):
        for assumeSafe in [True, False]:
            guess = {}
            guess[unrevealedCell] = assumeSafe
            # if assumeSafe:
            #     print("Guessing cell as safe:", unrevealedCell)
            # else:
            #     print("Guessing cell as mine:", unrevealedCell)

            solution = calcPermutationCounts(
                guess,
                set(unrevealedCell.neighbors),
                set(),
                permutationMap,
                totalMines - totalFlags - (0 if assumeSafe else 1),
                checkAnyValid=True,
            )
            if solution == 0:
                actions.append(
                    Action(
                        ActionType.FLAG if assumeSafe else ActionType.CLICK,
                        unrevealedCell.row,
                        unrevealedCell.col,
                    )
                )
                break
    return actions


def calcActionByProbability(cells: list[list[Cell]], totalMines: int) -> list[Action]:
    totalFlags = countFlagged(chain.from_iterable(cells))
    unreachedCells = filterUnreachedCells(cells)

    newCells = toCellExtended(cells)

    # permutate all possible combinations
    permutationCountsAll: dict[CellExtended, PermutationCount] = {}
    part = 0
    for unrevealedCell in filterUnrevealedCellExtended(newCells):
        if permutationCountsAll.get(unrevealedCell) is not None:  # already handled
            continue
        permutationCounts = defaultdict(lambda: PermutationCount(0, 0, part))
        calcPermutationCounts(
            {},
            set(unrevealedCell.neighbors),
            set(),
            permutationCounts,
            totalMines - totalFlags,
        )
        part += 1
        permutationCountsAll.update(permutationCounts)

    # calc unreached cells probability by total - known - expected mine from permutation
    if len(unreachedCells) > 0:
        expectMineCount = sum(c.mineProb() for c in permutationCountsAll.values())
        leavedMineCount = max(totalMines - totalFlags - expectMineCount, 0)
        prob = leavedMineCount / len(unreachedCells)
        for Cell in unreachedCells:
            cell = CellExtended(Cell.row, Cell.col, -1, [])
            permutationCountsAll[cell] = PermutationCount(
                mineCount=int(prob * 10000),
                safeCount=int((1 - prob) * 10000),
                part=part,
            )

    probSorted = sorted(
        permutationCountsAll.items(), key=lambda item: item[1].mineProb()
    )
    minProb = probSorted[0][1].mineProb()
    maxProb = probSorted[-1][1].mineProb()

    results: list[Action] = []

    if minProb < 0.01:
        cell, counts = probSorted[0]
        print(
            f"Found safe cell with probability {counts.mineProb():.2%} at ({cell.row}, {cell.col})"
        )
        results.append(
            Action(
                ActionType.CLICK,
                cell.row,
                cell.col,
            )
        )
    if maxProb > 0.99:
        cell, counts = probSorted[-1]
        print(
            f"Found mine cell with probability {counts.mineProb():.2%} at ({cell.row}, {cell.col})"
        )
        results.append(
            Action(
                ActionType.FLAG,
                cell.row,
                cell.col,
            )
        )

    if not results:
        print(f"Providing probability are between {minProb:.2%} and {maxProb:.2%}")
        targetProb = maxProb if minProb > 1 - maxProb else minProb
        for cell, counts in permutationCountsAll.items():
            isSuggested = counts.mineProb() == targetProb
            results.append(
                Action(
                    ActionType.HINT,
                    cell.row,
                    cell.col,
                    hint=f"{counts.mineProb():.0%}" + ("V" if isSuggested else ""),
                    hintDetail=f"Part {counts.part}: {counts.mineProb():.2%} ({counts.mineCount} / {counts.mineCount + counts.safeCount})",
                )
            )

    return results


def solve(cells: list[list[Cell]], totalMines: int) -> list[Action]:
    """
    Solve the minesweeper game using structured data.

    Args:
        cells: 2D list of Cell objects representing the game board

    Returns:
        List of Action objects to execute
    """
    actions = []

    if all(not cell.isRevealed for cell in chain.from_iterable(cells)):
        # First move, random click
        row = random.randint(0, len(cells) - 1)
        col = random.randint(0, len(cells[0]) - 1)
        actions.append(Action(ActionType.CLICK, row, col))
        return actions

    # Example: Basic solver logic
    for cell in chain.from_iterable(cells):
        neighbors = getNeighbors(cell, cells)
        unrevealed = filterUnrevealedNeighbors(neighbors)
        if not unrevealed:
            continue
        actions.extend(flagIfMatch(cell, neighbors))
        actions.extend(revealIfSafe(cell, neighbors))

    if not actions:
        print("Solve by false guessing")
        actions.extend(calcActionByFalseGuessing(cells, totalMines))
    if not actions:
        print("Solve by probability")
        actions.extend(calcActionByProbability(cells, totalMines))
    return actions


def example():
    board = [
        [-1, -1, -1],
        [1, 1, 1],
        [0, 0, 0],
    ]
    board = [
        [-1, -1, -1, -1],
        [2, 2, -1, -1],
        [1, -2, 2, 1],
    ]
    original_cells = [
        [
            Cell(row_idx, col_idx, val, val != -1, val == -2)
            for col_idx, val in enumerate(row)
        ]
        for row_idx, row in enumerate(board)
    ]
    print(solve(original_cells, 2))