Guesswork-free 3D minesweeper

minesweeper.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

from itertools import product, chain, count
import random, time

from z3 import *

N, M = 6, 11

def all_positions():
    return product(range(N), repeat=3)

def in_range(pos):
    return all(0 <= n < N for n in pos)

deltas = [(dx, dy, dz) for dx, dy, dz in product((-1, 0, 1), repeat=3) if not dx == dy == dz == 0]
def neighbours(pos):
    for d in deltas:
        new = tuple(x + dx for x, dx in zip(pos, d))
        if in_range(new):
            yield new

def generate_game(start_pos):
    def gen_pos():
        return tuple(random.randint(0, N-1) for _ in range(3))

    bombs = set()
    while len(bombs) < M:
        p = gen_pos()
        if p == start_pos or p in neighbours(start_pos):
            continue
        bombs.add(p)

    return bombs

def tag(bombs, pos):
    if pos in bombs:
        return -1
    return sum(nb in bombs for nb in neighbours(pos))

def is_uniquely_solvable(bombs, start):
    s = Solver()
    value = Function('value', IntSort(), IntSort(), IntSort(), IntSort())
    s.add(M == Sum([If(value(x, y, z) < 0, 1, 0) for x, y, z in all_positions()]))
    for pos in all_positions():
        nb_mines = Sum([If(value(*nb) < 0, 1, 0) for nb in neighbours(pos)])
        s.add(Or(value(*pos) == -1, value(*pos) == nb_mines))

    # All currently revealed safe positions
    known = set()
    # All positions that are known to be safe, but not currently revealed.
    to_visit = set([start])

    while to_visit:
        pos = to_visit.pop()
        # assert that the position we believe to be safe does not in fact contain a bomb...
        assert pos not in bombs
        t = tag(bombs, pos)
        # reveal the position
        known.add(pos)
        s.add(value(*pos) == t)

        if t == 0:
            to_visit.update(nb for nb in neighbours(pos) if nb not in known)
        elif len(to_visit) == 0 and len(known) < N**3:
            to_try = set(chain.from_iterable(neighbours(pos) for pos in known)) - known
            for cand in to_try:
                s.push()
                s.add(value(*cand) == -1)
                if s.check() == unsat:
                    s.pop()
                    to_visit.add(cand)
                    break
                s.pop()
            else:
                return False

    return True

def generate_solvable_game(start_pos):
    start = time.time()
    for tries in count(1):
        print('Generating solvable game [tries = {}]'.format(tries), end='\r')
        bombs = generate_game(start_pos)
        if is_uniquely_solvable(bombs, start_pos):
            print('Generated solvable game after {} tries ({:.3}s)'.format(tries, time.time() - start))
            return bombs

def print_grid(known):
    for z in range(N):
        print()
        print('Layer {}:'.format(z+1))
        for y in range(N):
            print(' '.join('{: ^3}'.format(known.get((x, y, z), '?')) for x in range(N)))


if __name__ == "__main__":
    def input_guess(prompt='Your guess? '):
        while True:
            try:
                x, y, z = map(int, input(prompt).split())
                if in_range((x, y, z)):
                    return (x, y, z)
            except ValueError:
                continue

    start_pos = input_guess('Initial guess? ')
    bombs = generate_solvable_game(start_pos)

    known = {}

    def reveal(pos):
        to_reveal = set([pos])
        while to_reveal:
            pos = to_reveal.pop()
            t = tag(bombs, pos)
            known[pos] = t
            if t == 0:
                to_reveal.update(nb for nb in neighbours(pos) if nb not in known)

    reveal(start_pos)

    while len(known) + M < N**3:
        print_grid(known)
        pos = input_guess()
        if pos in bombs:
            print('Dead!')
            break
        else:
            reveal(pos)

    print('You won!')