"""
 eight_queens_backtrack.py

 Search for a solution to the eight queens problem,
 in which the goal is put N queens on an NxN board
 in such a way that they can't attack each other.

 The approach here is a "backtrack" search, as described
 in Skienna chap 7. We add in pieces that head towards an
 answer, in a depth-first recursive search of the graph of
 possible queen placements.

   $ python eight_queens_backtrack.py -v

 See http://en.wikipedia.org/wiki/Eight_queens_puzzle for a
 much shorter python code using its built-in permutations library.

 Jim Mahoney | cs.marlboro.edu | April 1 2013 | MIT License
"""

class Board(object):
    """ An NxN chess board

        >>> a = Board(4)
        >>> a.board[1][3] = 'Q'
        >>> print a
           . . . .
           . . . Q
           . . . .
           . . . .
        >>> b = Board(4)
        >>> b.search()     # Print the two 4x4 solutions.
        --
           . Q . .
           . . . Q
           Q . . .
           . . Q .
        --
           . . Q .
           Q . . .
           . . . Q
           . Q . .
        >>> c = Board(8, verbose=False)
        >>> c.search()
        >>> print len(c.solutions)        # number of 8x8 solutions
        92
        >>> print c.solutions.keys()[0]   # first 8x8 solution
           . Q . . . . . .
           . . . . Q . . .
           . . . . . . Q .
           Q . . . . . . .
           . . Q . . . . .
           . . . . . . . Q
           . . . . . Q . .
           . . . Q . . . .

    """
    # Markers on the board are
    #    '.'    empty space
    #    'Q'    queen
    #    digit  place that can be attacked by queen on that row

    def __init__(self, N, verbose=True):
        self.N = N
        self.verbose = verbose  # if true, search() prints.
        self.board = [['.']*N for i in range(N)]
        self.solutions = {}

    def __str__(self):
        # Both 'x' or '.' print as '.'
        result = ''
        for row in self.board:
            result += '   ' + ' '.join(row) + '\n'
        import re
        return re.sub('\d+', '.', result).rstrip()

    def search(self, irow=0):
        """ Assuming queens on rows < irow are OK,
            put on each legit spot in this row, and recur.
            Print out the solutions found.
        """
        if irow >= self.N:                          # solution?
            self.solutions[str(self)] = True        #   remember it
            if self.verbose:                        #   output it
                print '--'
                print self
        else:
            for icol in range(self.N):
                if self.board[irow][icol] == '.':   # can continue?
                    self.put_queen(irow, icol)      #   setup
                    self.search(irow + 1)           #   recur
                    self.remove_queen(irow, icol)   #   undo setup

    def put_queen(self, irow, icol):
        """ Put a queen at the given spot, and mark
            the squares below that it can attack with it's row digit"""
        self.change_mark(irow, icol, '.', 'Q')
        irow_below = irow + 1
        icol_diag_left = icol - 1
        icol_diag_right = icol + 1
        while irow_below < self.N:
            self.change_mark(irow_below, icol,            '.', irow)
            self.change_mark(irow_below, icol_diag_right, '.', irow)
            self.change_mark(irow_below, icol_diag_left,  '.', irow)            
            irow_below += 1
            icol_diag_left -= 1
            icol_diag_right += 1

    def remove_queen(self, irow, icol):
        """ Remove marker for queen at irow, icol,
            and remove marks for spaces it attacks below it."""
        self.change_mark(irow, icol, 'Q', '.')
        irow_below = irow + 1
        icol_diag_left = icol - 1
        icol_diag_right = icol + 1
        while irow_below < self.N:
            self.change_mark(irow_below, icol,            irow, '.')
            self.change_mark(irow_below, icol_diag_right, irow, '.')
            self.change_mark(irow_below, icol_diag_left,  irow, '.')
            irow_below += 1
            icol_diag_left -= 1
            icol_diag_right += 1

    def change_mark(self, irow, icol, old_mark, new_mark):
        """ put new_mark at irow,icol if old_mark was there """
        old_mark = str(old_mark)
        new_mark = str(new_mark)
        if icol >= 0 and icol < self.N and self.board[irow][icol] == old_mark:
                self.board[irow][icol] = new_mark

if __name__ == "__main__":
    import doctest
    doctest.testmod()