"""
 utilities.py

 utility functions for n-queens problems

 The board representation is a list of integers from 0 to n-1,
 where each integer gives column that a queen is in for that row.
 For example, the n=4 board [0, 3, 1, 1] would be
      * . . .
      . . . *
      . * . .
      . * . .
 where * is a queen and . is an empty square, which is *not*
 a solution because (a) there are two queens in coolumn 1,
 and (b) the queens in row 1 and row 3 are on the same diagonal.

 Jim Mahoney | cs.bennington.college | May 2022 | MIT License
"""

from random import randrange

valid_board_5 = [1, 3, 0, 2, 4]    # an example n=5 solution
#    0 1 2 3 4
# 0  . * . . .    queen in column 1    
# 1  . . . * .                    3
# 2  * . . . .                    0
# 3  . . * . .                    2
# 4  . . . . *                    4

invalid_diagonal_5 = [3, 1, 0, 2, 4]  # an n=5 board with two invalid diagonals
#    0 1 2 3 4
# 0  . . . * .    queen in column 3
# 1  . * . . .                    1    <--+  <--+  same diagonal
# 2  * . . . .                    0       |  <--+
# 3  . . * . .                    2       |
# 4  . . . . *                    4    <--+

invalid_column_5 = [3, 1, 0, 2, 4]  # an n=5 board with queens in the same column
#    0 1 2 3 4
# 0  . * . . .    queen in column 1   <--+  same column  <--+ same diagonal
# 1  . . . * .                    3      |                  |
# 2  * . . . .                    0      |                  |
# 3  . . . . *                    5      |               <--+
# 4  . * . . .                    1   <--+       

def board_string(board):
    """ return a printable string given a board """
    result = "\n"
    n = len(board)
    for i in range(n):
        result += " "
        for j in range(n):
            if board[i] == j:
                result += "* "
            else:
                result += ". "
        result += "\n"
    return result

def random_board(n):
    """ return a representation of a random board of size n """
    # One queen per row; some are likely on same column and/or diagonal.
    # The syntax here uses python's list comprehsnsion; see
    # for example https://docs.python.org/3/tutorial/datastructures.html
    return [randrange(0, n) for row in range(n)]

def same_diagonal(board, row1, row2):
    """ True if the queens in row1, row2 are on the same diagonal """
    # Two queens are on the same diagonal if their horizontal distance
    # apart (i-j) is the same as their vertical distance
    # which is abs(board[i]-board[j]).
    return abs(row1 - row2) == abs(board[row1] - board[row2])

def same_column(board, row1, row2):
    """ True if the queens in row1, row2 are on the same column """
    return board[row1] == board[row2]

def invalid_below(board, row):
    """ True if queen in this row is on a digonal or column  with a lower row """
    for row1 in range(row):
        if same_diagonal(board, row1, row) or same_column(board, row1, row):
            return True
    return False
        
def valid(board):
    """ True if the board has no two queens on a  diagonal or column """
    for row in range(1, len(board)):
        if invalid_below(board, row):
            return False
    return True

def attacking_pairs(board):
    """ Return the number of pairs of queens which can attack each other """
    attacks = 0
    n = len(board)
    for row1 in range(n):
        for row2 in range(row1):
            if board[row1] == board[row2]:
                attacks += 1
            if same_diagonal(board, row1, row2):
                attacks += 1
    return attacks

# --- tests ---

assert valid(valid_board_5), 'valid on good board is OK'
assert not valid(invalid_diagonal_5), 'valid on bad board is OK'
assert board_string(valid_board_5) == """
 . * . . . 
 . . . * . 
 * . . . . 
 . . * . . 
 . . . . * 
""", 'board_string OK'
assert attacking_pairs(valid_board_5) == 0, "attacking_pairs 0"
assert attacking_pairs(invalid_diagonal_5) == 2, "attacking_pairs diagonal"
assert attacking_pairs(invalid_column_5) == 2, "attacking_pairs col & diag"

if __name__ == '__main__':
    print("Passed all tests.")