#lang scheme/base
(require
 scheme/match
 (only-in (lib "13.ss" "srfi") string-join))

;; Eight queens problem, using continuation passing style
;; to handle the backtracking.
;;
;; This example uses some mutation, so the backtracking
;; is nontrivial since it has to undo what it did before.


;; Let's represent a board as a square matrix of booleans.
;; Internally, we're using a linear representation.
;; A queen is on a place when the value is true.
(define-struct board (vals n) #:transparent)
(define (new-board n)
  (make-board (make-vector (square n) #f) n))


;; square: number -> number
;; Your standard squaring function.
(define (square x) (* x x))


;; board-copy: board -> board
;; Returns a new board initialized with the contents of
;; the input board.
(define (board-copy a-board)
  (make-board (build-vector (square (board-n a-board))
                            (lambda (i)
                              (vector-ref (board-vals a-board) i)))
              (board-n a-board)))


;; board-ref: board number number -> boolean
;; Returns the element at the given row/col.
(define (board-ref a-board a-row a-col)
  (match a-board
    [(struct board (vals n))
     (vector-ref vals (+ (* n a-row) a-col))]))


;; board-set!: board number number boolean -> void
;; Sets the value at the given row/col.
(define (board-set! a-board a-row a-col v)
  (match a-board
    [(struct board (vals n))
     (vector-set! vals (+ (* n a-row) a-col) v)]))


;; board->string: board -> string
;; Stringifies a board to make it easier to see what's's happening.
(define (board->string a-board)
  (let ([rows
         (for/list ([i (in-range (board-n a-board))])
           (apply string-append
                  (for/list ([j (in-range (board-n a-board))])
                    (cond
                      [(board-ref a-board i j)
                       "Q"]
                      [else
                       "."]))))])
    (string-join rows "\n")))


;; start-search: board success-continuation failure-contination -> any
;; Begins a search for solutions.  Warning: this search will munge
;; up the given a-board.
(define (start-search! a-board succ-k fail-k)
  (search! a-board (board-n a-board) 0 succ-k fail-k))


;; queen-search: board number number (board (-> any) -> any) (-> any) -> any
;; Tries to place n-queens on the board.
;;
;; We direct the search to find solutions more quickly by
;; placing a precondition: we assume that all rows before next-row
;; are occupied, and that all rows starting from next-row are
;; unoccupied.
;;
;; If a candidate is successful, we call the success continuation and pass it
;; a thunk that will continue the search for candidates.
(define (search! a-board n-queens next-row succ-k fail-k)
  (cond
    [(= n-queens 0)
     (succ-k a-board fail-k)]
    [else
     (place-a-queen-and-try-again! a-board n-queens next-row succ-k fail-k)]))


;; place-a-queen-and-try-again: board number number succ-cont fail-cont
;; Tries to place a queen at row i, on some unoccupied column, and
;; continue the search from there.
(define (place-a-queen-and-try-again! a-board n-queens a-row
                                     succ-k fail-k)
  (let loop ([a-col 0])
    (cond
      [(= a-col (board-n a-board))
       (fail-k)]
      [(and (board-diagonals-clear? a-board a-row a-col)
            (board-col-clear? a-board a-col))
       (board-set! a-board a-row a-col #t)
       (search! a-board
                (sub1 n-queens)
                (add1 a-row)
                succ-k
                (lambda ()
                  (board-set! a-board a-row a-col #f)
                  (loop (add1 a-col))))]
      [else
       (loop (add1 a-col))])))


;; board-col-clear?: board number -> boolean
;; Returns true if the col's booleans are all #f.
(define (board-col-clear? a-board a-col)
  (match a-board
    [(struct board (rows n))
     (not
      (for/or ([a-row (in-range n)])
              (board-ref a-board a-row a-col)))]))


;; board-diagonals-clear?: board number number -> boolean
;; Returns true if the diagonals are clear of queens.
(define (board-diagonals-clear? a-board a-row a-col)
  (define (good-index? r)
    (<= 0 r (sub1 (board-n a-board))))
  
  (define (quadrant-clear? row-plus-or-minus col-plus-or-minus)
    (let loop ([i 0])
      (cond
        [(or (not (good-index? (row-plus-or-minus a-row i)))
             (not (good-index? (col-plus-or-minus a-col i))))
         #t]
        [(board-ref a-board
                    (row-plus-or-minus a-row i)
                    (col-plus-or-minus a-col i))
         #f]
        [else
         (loop (add1 i))])))
  (and (quadrant-clear? + +)
       (quadrant-clear? + -)
       (quadrant-clear? - +)
       (quadrant-clear? - -)))


;; print-solutions: number -> void
;; Prints out the solutions as we find them.
(define (print-solutions n)
  (start-search! (new-board n)
                (lambda (solution try-again)
                  (printf "~a~n~n" (board->string solution))
                  (try-again))
                (lambda ()
                  (printf "Done!~n"))))


;; get-solutions: number -> (listof board)
;; Collects a list of solutions.
(define (get-solutions n)
  (start-search! (new-board n)
                (lambda (solution try-again)
                  (let ([new-solution (board-copy solution)])
                    (cons new-solution (try-again))))
                (lambda ()
                  '())))


;; count-solutions: number -> number
;; Returns the number of unique solutions.
(define (count-solutions n)
  (start-search! (new-board n)
                 (lambda (solution try-again)
                   (add1 (try-again)))
                 (lambda ()
                   0)))