#!/usr/bin/env python
##
## Name:     sudoku.py
## Purpose:  Generate and solve Sudoku puzzles.
## Author:   M. J. Fromberger <http://www.dartmouth.edu/~sting/>
## Info:     $Id: sudoku.py 320 2007-08-01 14:27:39Z sting $
##

import math, random, re

# {{ class ConsistencyError 

class ConsistencyError (Exception):
    """Exception raised during puzzle solution when a consistency
    error is detected.
    """
    pass

# }}

# {{ class Puzzle 

class Puzzle (object):
    """Represents a square Sudoku puzzle with rectangular subregions.
    
    Subclass and override ._getbox() to implement more complex shapes
    for subregions, such as jigsaw Sudoku.
    """
    def __init__(self, dim = 9, hdiv = 3, init = None):
        """Initialize a new blank square Sudoku board with dim rows
        and columns, and hdiv-width rectangular subdivisions.  It is
        required that hdiv be a divisor of dim.

        If init is not None, it is passed to the .load() method to
        initialize the contents of the board; otherwise the board
        begins blank.
        """
        assert dim > 0
        assert math.modf(dim / hdiv)[0] == 0.0
        
        self._dim = dim
        self._hdiv = hdiv
        self._vdiv = dim / hdiv
        if init is None:
            self.reset()
        else:
            self.load(init)

    def reset(self):
        """Erase the board to blank."""
        nr = nc = self._dim
        
        self._brd = [None] * nr
        for row in xrange(nr):
            self._brd[row] = [None] * nc
            for col in xrange(nc):
                self._brd[row][col] = set(xrange(1, self._dim + 1))

    def nrows(self):
        """Returns the dimensions of the board."""
        return self._dim

    ncols = nrows
    
    def rows(self):
        """Return all the rows of the board, top to bottom."""
        return iter(self._getrow(r)
                    for r in xrange(self.nrows()))

    def cols(self):
        """Return all the columns of the board, left to right."""
        return iter(self._getcol(c)
                    for c in xrange(self.ncols()))

    def load(self, p):
        """Load a puzzle into this board.  The puzzle should be given
        as a string of numbers of the appropriate dimensions.  White
        space delimits values, and "_" is used to mark a blank square.

        Raises ValueError if the string is of the wrong dimensions, or
        contains any unexpected characters.
        """
        q = re.split(r'\s+', p.strip())
        nr = self.nrows()
        nc = self.ncols()
        if len(q) <> (nr * nc):
            raise ValueError("Incorrect puzzle dimension")
        
        self.reset()
        
        pos = 0
        for ch in q:
            if ch == '_':
                pass
            else:
                r = pos // nc
                c = pos % nc
                try:
                    v = int(ch)
                except ValueError:
                    raise ValueError("Invalid puzzle entry: %s" % ch)
                if v < 1 or v > self._dim:
                    raise ValueError("Value out of range: %s" % v)

                s = self._brd[r][c]
                s.clear()
                s.add(v)
            
            pos += 1

    def copy(self):
        """Return a copy of the puzzle as it is currently."""
        out = type(self)(self._dim, self._hdiv)
        out._brd = self._copyboard()
        return out
    
    def unsolved(self):
        """Find an unsolved square of the board, returning its row and
        column index.  Returns None, None if no squares are unsolved.
        """
        cr = cc = None
        cv = self._dim * 2

        for r in xrange(self.nrows()):
            for c in xrange(self.ncols()):
                sv = len(self._getcel(r, c))
                
                if sv > 1 and sv < cv:
                    cr, cc, cv = r, c, sv
        
        return cr, cc
    
    def solve(self):
        """Attempt to fill in a solution for the puzzle.  If a
        solution is found, returns True; otherwise returns False.

        Note: If you call .solve() on a blank board, it generates a
        complete board, chosen at random.  If there is more than one
        possible solution to the current board state, only one such
        state is actually found.
        """
        self._count = 0
        
        # Check for already-solved clues
        for r in xrange(self.nrows()):
            for c in xrange(self.ncols()):
                s = self._getcel(r, c)
                if len(s) == 1:
                    self._setcel(r, c, list(s)[0])

        # Initialize the work queue; this contains the backtrack
        # states for a depth-first scan of the unsolved spaces in the
        # board.
        r, c = self.unsolved()
        if r is None:
            return True
        
        more = [(r, c, self._brd, v)
                for v in self._getcel(r, c)]
        random.shuffle(more)

        # At the start of this loop, the board is unsolved and "more"
        # contains the possible next steps toward a solution; if it is
        # empty, there is no viable solution.
        
        while more:
            self._count += 1
            r, c, s, v = more.pop()
            self._brd  = s
            try:
                self._setcel(r, c, v) # may fail

                r, c = self.unsolved()
                if r is None:
                    return True

                fwd = list((r, c, self._brd, v)
                           for v in self._getcel(r, c))
                random.shuffle(fwd)
                more.extend(fwd)
            except ConsistencyError:
                pass

        return False # no possible solution
    
    def _getcel(self, rx, cx):
        """[private] Get the contents of the specified square."""
        return self._brd[rx][cx]

    def _clearcel(self, rx, cx):
        """[private] Clear the specified square."""
        s = self._getcel(rx, cx)
        s.clear()
        s.update(xrange(1, self._dim + 1))
    
    def _setcel(self, rx, cx, v):
        """[private] Set the contents of the specified square to v,
        and propagate its implications to the containing row, column,
        and box of that square.  Raises ConsistencyError if there is
        a problem, and restores the original board state.
        """
        save = self._brd
        try:
            self._brd = self._copyboard()

            s = self._getcel(rx, cx)
            s.clear()
            s.add(v)
            
            self._checkcel(rx, cx, v)
        except ConsistencyError:
            self._brd = save
            raise

    def _checkcel(self, rx, cx, v):
        """[private] Propagate implications of a value change to the
        row, column, and box neighbors of a square.

        Raises ConsistencyError if there is a problem.
        """
        fix = set()

        for c, s in enumerate(self._getrow(rx)):
            os = len(s)
            if c <> cx:
                s.discard(v)
                if os > 1 and len(s) == 1:
                    fix.add((rx, c, list(s)[0]))
                elif len(s) == 0:
                    raise ConsistencyError(rx, c)
        for r, s in enumerate(self._getcol(cx)):
            os = len(s)
            if r <> rx:
                s.discard(v)
                if os > 1 and len(s) == 1:
                    fix.add((r, cx, list(s)[0]))
                elif len(s) == 0:
                    raise ConsistencyError(r, cx)
        for r, c, s in self._getbox(rx, cx):
            os = len(s)
            if r <> rx or c <> cx:
                s.discard(v)
                if os > 1 and len(s) == 1:
                    fix.add((r, c, list(s)[0]))
                elif len(s) == 0:
                    raise ConsistencyError(r, c)

        for r, c, v in fix:
            self._checkcel(r, c, v)
    
    def _getrow(self, rx):
        """[private] Get the row whose index is given."""
        return self._brd[rx]

    def _getcol(self, cx):
        """[private] Get the column whose index is given."""
        return list(self._getcel(rx, cx)
                    for rx in xrange(self.nrows()))

    def _getbox(self, rx, cx):
        """[private] Get the box containing the square whose row and
        column indices are given.  This is returned as a list of
        tuples (r, c, s) where r, c are the coordinates of a square
        and s is the content set for that square.
        """
        br = (rx // self._hdiv) * self._hdiv
        bc = (cx // self._vdiv) * self._vdiv

        out = [None] * (self._hdiv * self._vdiv)
        pos = 0
        for r in xrange(br, br + self._hdiv):
            for c in xrange(bc, bc + self._vdiv):
                out[pos] = r, c, self._getcel(r, c)
                pos += 1
        return out

    def _copyboard(self):
        """[private] Make a detached copy of the board contents."""
        return [[set(s) for s in row]
                for row in self._brd]
    
    def as_text(self):
        """Renders the board as human-readable text."""
        out = ''
        ln  = len(str(self._dim))
        fmt = ' %%%ds ' % ln
        for rx, row in enumerate(self.rows()):
            out += '|'.join(fmt % setrep(s)
                            for s in row)
            if rx < self._dim - 1:
                if rx % self._vdiv == (self._vdiv - 1):
                    t = '=' * (ln + 2)
                else:
                    t = '-' * (ln + 2)
                out += '\n'
                out += '+'.join(t for s in row)
                out += '\n'
        return out
    
    def __str__(self):
        ln = len(str(self._dim))
        fmt = '%%%ds' % ln
        out = '\n'.join(' '.join(fmt % setrep(s, '_') for s in row)
                        for row in self.rows())
        return out

    def __eq__(self, other):
        """Compare two puzzle boards; they are equal iff they have the
        same dimensions and equal content.
        """
        if not isinstance(other, type(self)):
            return False
        else:
            return str(self) == str(other)

# }}

# {{ class JigsawPuzzle 

class JigsawPuzzle (Puzzle):
    """A variant of the Sudoku puzzle that allows for arbitrary
    container regions instead of square boxes.  Subclass and supply
    a value for the REGIONS member.

    To specify regions, supply a string of characters separated by
    whitespace.  One nonspace character must be supplied for each cell
    of the puzzle, in row-major order.  All those cells which share a
    region will be considered to be in the same "box".  In order to be
    legal, each region must have the same number of elements, and that
    number must match the dimension of the puzzle.
    """
    def __init__(self, dim = 9, init = None):
        super(JigsawPuzzle, self).__init__(dim, dim, init)
        self._checkreg()
    
    def _checkreg(self):
        """[private] Check the REGIONS field and build a quick lookup
        table to match positions to region contents.
        """
        try:
            reg = self.REGIONS
        except AttributeError:
            raise TypeError("No value for REGIONS supplied")
        
        reg  = re.split(r'\s+', reg.strip())
        
        if len(reg) <> (self._dim ** 2):
            raise TypeError("Improper REGIONS value supplied")

        pos = 0 ; rmap = {}
        for r in xrange(self.nrows()):
            for c in xrange(self.ncols()):
                rmap[r, c] = reg[pos]
                pos += 1

        cmap = dict((v, set()) for v in set(reg))
        for k, v in rmap.iteritems():
            cmap[v].add(k)
        
        for k, v in cmap.iteritems():
            if len(v) <> self._dim:
                raise TypeError("Improper region size %s, %d" %
                                (k, len(v)))

        self._rmap = rmap
        self._cmap = cmap
    
    def _getbox(self, r, c):
        """[private] Return the members of the region containing the
        given cell.  Overrides the inherited version from Puzzle.
        """
        tag = self._rmap[r, c]
        pts = self._cmap[tag]

        return list(v + (self._brd[r][c],)
                    for v in pts)

# }}

# {{ setrep(s[, alt]) 

def setrep(s, alt = ' '):
    """[private] Generate a human-readable representation of a cell.
    """
    if len(s) == 1:
        return str(list(s)[0])
    else:
        return alt

# }}

# {{ test() 

def test():
    """Unit test for Sudoku solver."""
    global puzzle1, puzzle1s
    global puzzle2, puzzle2s
    global puzzle3, puzzle3s
    
    puzzle1 = '''
1 2 6 _ 9 _ _ _ _
_ _ 5 7 _ _ _ 1 3
_ _ _ 6 _ _ _ _ _
4 _ _ 5 8 _ 7 _ 9
3 5 _ 4 _ _ _ _ 1
_ _ _ 2 7 _ 4 _ _
_ 6 3 8 _ _ _ 9 _
_ _ 8 _ _ _ 5 7 _
_ _ _ _ 4 6 _ 2 _
'''
    puzzle1s = '''
1 2 6 3 9 5 8 4 7
8 9 5 7 2 4 6 1 3
7 3 4 6 1 8 9 5 2
4 1 2 5 8 3 7 6 9
3 5 7 4 6 9 2 8 1
6 8 9 2 7 1 4 3 5
2 6 3 8 5 7 1 9 4
9 4 8 1 3 2 5 7 6
5 7 1 9 4 6 3 2 8
'''
    puzzle2 = '''
2 _ 5 _ _ 6 _ _ _
1 _ _ _ _ 3 _ _ 8
6 3 _ 5 4 7 _ _ 1
_ _ _ 1 _ _ 6 _ _
_ _ _ _ 3 4 _ 8 _
_ 8 _ _ 9 _ _ 7 _
_ 7 2 _ _ _ _ _ _
_ _ 6 _ 5 9 _ 2 3
8 _ _ _ _ 2 _ _ _
'''
    puzzle2s = '''
2 9 5 8 1 6 3 4 7
1 4 7 9 2 3 5 6 8
6 3 8 5 4 7 2 9 1
5 2 9 1 7 8 6 3 4
7 6 1 2 3 4 9 8 5
3 8 4 6 9 5 1 7 2
9 7 2 3 8 1 4 5 6
4 1 6 7 5 9 8 2 3
8 5 3 4 6 2 7 1 9
'''
    # This is a very difficult puzzle to challenge brute-force
    # Sudoku solvers.  It is linked from:
    #  <http://en.wikipedia.org/wiki/Algorithmics_of_sudoku>
    puzzle3 = '''
_ _ _ _ _ _ _ _ _
_ _ _ _ _ 3 _ 8 5
_ _ 1 _ 2 _ _ _ _
_ _ _ 5 _ 7 _ _ _
_ _ 4 _ _ _ 1 _ _
_ 9 _ _ _ _ _ _ _
5 _ _ _ _ _ _ 7 3
_ _ 2 _ 1 _ _ _ _
_ _ _ _ 4 _ _ _ 9
'''
    puzzle3s = '''
9 8 7 6 5 4 3 2 1
2 4 6 1 7 3 9 8 5
3 5 1 9 2 8 7 4 6
1 2 8 5 3 7 6 9 4
6 3 4 8 9 2 1 5 7
7 9 5 4 6 1 8 3 2
5 1 9 2 8 6 4 7 3
4 7 2 3 1 9 5 6 8
8 6 3 7 4 5 2 1 9
'''
    p1  = Puzzle(init = puzzle1)
    p1s = Puzzle(init = puzzle1s)
    p2  = Puzzle(init = puzzle2)
    p2s = Puzzle(init = puzzle2s)
    p3  = Puzzle(init = puzzle3)
    p3s = Puzzle(init = puzzle3s)
    
    assert p1.solve() == True
    if p1 == p1s:
        print "p1 solved, %d iterations" % p1._count
    else:
        print "p1 failed, got\n%s\nexpected\n%s" % (p1, p1s)

    assert p2.solve() == True
    if p2 == p2s:
        print "p2 solved, %d iterations" % p2._count
    else:
        print "p2 failed, got\n%s\nexpected\n%s" % (p2, p2s)

    assert p3.solve() == True
    if p3 == p3s:
        print "p3 solved, %d iterations" % p3._count
    else:
        print "p3 failed, got\n%s\nexpected\n%s" % (p3, p3s)

# }}

__all__ = ('Puzzle', 'JigsawPuzzle')

# Here there be dragons