## This program finds arrangements of q queens and p pawns on an r-by-c chessboard (q,p,r, and c input by user) so that any pair of queens on the same row, column, or diagonal has at least one pawn between them.
## The program is inspired by the September 2018 College Mathematics Journal article "It's Puzzling" by C. Douglas Howard.  (https://www.tandfonline.com/doi/abs/10.1080/07468342.2018.1500062?journalCode=ucmj20)

## Written in Python 3.6.4

## Doug Chatham
## d.chatham@moreheadstate.edu
## July 14, 2019

from random import *
from math import *
from time import *

def Attack(pieces,i,j,p,q):
# Does piece i attack piece j?  Yes if they are on same row, column, rising diagonal, or falling diagonal and there is no pawn between them.
# Otherwise, no.

    if pieces[i][0]==pieces[j][0]:
        for k in range(q,q+p):
            if pieces[k][0]==pieces[i][0]:
                if (min(pieces[i][1], pieces[j][1])<=pieces[k][1]) and (pieces[k][1] <= max(pieces[i][1], pieces[j][1])):
                                                                        return False
        return True                                                                
    if pieces[i][1]==pieces[j][1]:
        for k in range(q,q+p):
            if pieces[k][1]==pieces[i][1]:
                if (min(pieces[i][0], pieces[j][0])<=pieces[k][0]) and (pieces[k][0] <= max(pieces[i][0], pieces[j][0])):
                                                                        return False
        return True                                                                  
    if pieces[i][0]+pieces[i][1]==pieces[j][0]+pieces[j][1]:
        for k in range(q,q+p):
            if pieces[k][0]+pieces[k][1]==pieces[i][0]+pieces[i][1]:
                if (min(pieces[i][0], pieces[j][0])<=pieces[k][0]) and (pieces[k][0] <= max(pieces[i][0], pieces[j][0])):
                                                                        return False
        return True  
    if pieces[i][0]-pieces[i][1]==pieces[j][0]-pieces[j][1]:
        for k in range(q,q+p):
            if pieces[k][0]-pieces[k][1]==pieces[i][0]-pieces[i][1]:
                if (min(pieces[i][0], pieces[j][0])<=pieces[k][0]) and (pieces[k][0] <= max(pieces[i][0], pieces[j][0])):
                                                                        return False
        return True  
    return False
    
def Energy(pieces,p,q):
# The energy of an arrangement of queens and pawns is the number of pairs of queens that attack each other.

    tot=0
    for i in range(q):
        for j in range(i+1,q):
            if Attack(pieces,i,j,p,q):
                tot=tot+1
    return tot            

def Move(pieces,p,q,r,c):
# Randomly select a piece to move and a proposed position to move it to.

    newpieces=[[0,0] for i in range(q+p)]
    changerow=randrange(r)
    changecol=randrange(c)
    if q+p<0.75*r*c: #if board not nearly full of pieces, require target square to be empty
        while ([changerow,changecol] in pieces):
            changerow=randrange(r)
            changecol=randrange(c)
    changepos=randrange(q+p)
    newpieces[changepos][0]=changerow
    newpieces[changepos][1]=changecol
    for i in range(q+p):
        if i!=changepos:
            newpieces[i][0]=pieces[i][0]
            newpieces[i][1]=pieces[i][1]
    if ([changerow,changecol] in pieces): #if target square was not empty, put its piece in the originating square
        oldpos=0
        for i in range(q+p):
            if pieces[i][0]==changerow and pieces[i][1]==changecol:
                oldpos=i
        newpieces[oldpos][0]=pieces[changepos][0]
        newpieces[oldpos][1]=pieces[changepos][1]
    return newpieces

def Pick(p,q,r,c):
# Randomly determine an initial arrangement of q queens and p pawns on an r-by-c chessboard.
# Note that pieces[0:q] will contain the positions of the q queens and pieces[q:q+p] will contain the positions of the p pawns.

    pieces=[[0,0] for i in range(q+p)]
    sel=sample(range(r*c),q+p)
    for i in range(q+p):
       pieces[i][0]=int(sel[i]/c)
       pieces[i][1]=sel[i]-c*pieces[i][0]
    return pieces

def Display(pieces,r,c,p,q):
# Output an arrangement in a 2D board display

    for i in range(r):
        for j in range(c):
            if [i,j] in pieces[0:q]:
                print("Q",end="")
            else:
                if [i,j] in pieces[q:p+q]:
                    print("P",end="")
                else:
                    print(".",end="")
        print(" ",end='\n')
    print(" ",end='\n')
    return

# Main program

a=seed()
temp=input("Temperature:") # a parameter that controls the probability of moving to a higher-energy state. 
temp=float(temp)
r=input("Number of rows:")
r=int(r)
c=input("Number of columns:")
c=int(c)
q=input("Number of queens:")
q=int(q)
p=input("Number of pawns:")
p=int(p)
mul=input("Number of solutions:") # Actually, number of times to *attempt* to find a solution
mul=int(mul)
maxtime=input("Maximum time per solution:") # Enter 0 to have no limit on search time.
maxtime=float(maxtime)
for index in range(mul):
    t0=clock()
    paces = 1
    looks = 1
    pieces = Pick(p,q,r,c)
    oldEng = Energy(pieces,p,q)
    while oldEng>0 and (maxtime==0 or clock()-t0<maxtime): # if oldEng == 0, we have a solution
        newpieces=Move(pieces,p,q,r,c)
        newEng= Energy(newpieces,p,q)
        diffEng=newEng-oldEng
        looks = looks + 1
        if diffEng<0 or random()<exp(-diffEng/temp): # if a proposed change lowers the energy or a randomly chosen float number between 0 and 1 is less than exp(-diffEng/temp), then replace the old arrrangement with the proposed one.
            pieces=newpieces
            oldEng=newEng
            paces = paces + 1
    if oldEng>0:
        print("Energy:",oldEng)
    print("\nTime:",clock()-t0," seconds")
    print("\nMoves considered:",looks)
    print("\nSteps taken:",paces)
    print("\nQueens:",sorted(pieces[0:q]))
    print("Pawns:",sorted(pieces[q:q+p]),"\n")
    Display(pieces,r,c,p,q)
y=input("Hit enter to end program:")

                                                                        
