#!/usr/bin/python3

# A solver for the n-queens problem
#
# It uses a recursive search algorithm with backtracking,
# and displays naively the current partial solution.
#
# Some background information on Wikipedia:
#   https://en.wikipedia.org/wiki/Eight_queens_puzzle
#
# Question 1
# ----------
# Modify the code to add a more graphical display of the current
# partial solution, using pygame [1] or curses [2].
# The "old" display should remain available and the search
# algorithm should not be duplicated. More specifically,
# search(size,G) should start the search with a graphical progress
# display, and start(size,O) should do the same with the old
# display, for some G and O. The two modes should be available to
# the user via the command-line.
#
# [1] https://docs.python.org/3/howto/curses.html
# [2] https://www.pygame.org/
#
# Question 2
# ----------
# Modify the procedure so that it also allows to obtain
# the number of solutions instead of printing and exiting
# after having found the first solution.
# It should still be possible to obtain the old behaviour
# (which avoids exploring the whole solution space).
# The two modes should be available to the user via the
# command-line.
# Bonus points if other interesting uses are enabled by the
# modified procedure.

def search(size):
    # The algorithm attempts to set one queen per line,
    # starting with line 0 and processing lines in order,
    # backtracking when it stuck.
    #
    # In order to test rapidly for possible positions, it
    # maintains tables of free columns and diagonals:
    #  - col[j] indicates that column j is free
    #  - up[k] indicates that the upward-going diagonal k is free
    #  - up[k] indicates that the downward-going diagonal k is free
    # Here j belongs to range(size) and k belong to range(2*size-1).
    col = [-1 for _ in range(size)]
    up = [-1 for _ in range(2*size-1)]
    down = [-1 for _ in range(2*size-1)]
    # The current solution is maintained for display as a list of
    # successive column indices.
    sol = []
    def searchline(i):
        if i==size:
            print("Solution found.")
            exit(0)
        for j in range(size):
            if col[j] and down[i+j] and up[i-j+size-1]:
                col[j] = down[i+j] = up [i-j+size-1] = False
                sol.append(j)
                print(sol)
                searchline(i+1)
                sol.pop()
                col[j] = down[i+j] = up [i-j+size-1] = True
    searchline(0)
    print("No solution!")
    exit(0)

import sys

if __name__ == "__main__":
    size = 7
    if len(sys.argv)>1:
        try:
            size = int(sys.argv[1])
        except (IndexError,ValueError):
            print("Usage: %s <int>" % sys.argv[0])
            exit(1)
    print("Searching for size %d..." % size)
    search(size)