# ENGG 2440A / ESTR 2004 Lecture 1 -- python code
# Requires: python satispy package, minisat solver installed

from satispy import Variable, Cnf
from satispy.solver import Minisat

# Place N queens on an N by N chessboard so that no two attack one another
def queens(N = 8):
    # N * N array of variables P_{ij} indicating the presence of a queen in square [i][j]
    queenat = [[Variable("P_{" + str(i) + str(j) + "}") for i in range(N)] for j in range(N)]
    
    # propositional formula describing the assignment of queens to the board
    formula = Cnf()
    
    # at least one queen in every row
    for i in range(N):
        oneinthisrow = Cnf()
        for j in range(N):
            oneinthisrow |= queenat[i][j]
        formula &= oneinthisrow
        
    # at least one queen in every column
    for j in range(N):
        oneinthiscol = Cnf()
        for i in range(N):
            oneinthiscol |= queenat[i][j]
        formula &= oneinthiscol
    
    # disallow pairs that live in same row, same column, or same diagonal
    for i1 in range(N):
        for j1 in range(N):
            for i2 in range(N):
                for j2 in range(N):
                    # break the symmetry between the pairs [i1][j1] and [i2][j2]
                    if N * i1 + j1 < N * i2 + j2:
                        # check if squares [i1][j1] and [i2][j2] share a row, column, or diagonal            
                        if (i1 == i2) or (j1 == j2) or (abs(i2 - i1) == abs(j2 - j1)):
                            formula &= (-queenat[i1][j1]) | (-queenat[i2][j2])
                        
    # solve the instance using the minisat solver
    solution = Minisat().solve(formula)
    
    if solution.error != False:
        print "Error: " + solution.error
    elif solution.success:
        # construct and print the chessboard
        chessboard = ""
        for j in range(N):
            for i in range(N):
                if solution[queenat[i][j]]:
                    chessboard += "Q"
                else:
                    chessboard += "."
            chessboard += "\n"
        print chessboard
    else:
        print "The queens cannot be placed on the chessboard"