#lang scheme

;; O(n^n) sudoku solver and board generator
;; provides
;;  'create-sudoku' num num :: board
;;  'solve' board :: board
;;  'print-board' board :: void
;;  'solved?' board :: boolean
;;
;; Input should be piped to stdin.
;; cat f | mzscheme sudoku.ss
;;
;; Format should input should be like the following
;;
;; 2 3
;; 1 2 _ 4 5 6
;; 3 4 _ 6 1 2
;; 5 6 _ 2 3 4
;; 2 3 4 5 _ 1
;; 4 5 6 1 _ 3
;; 6 1 2 3 _ 5

(require mzlib/pregexp)

;; cells is list-of-list-of-num num num
(define-struct board (cells width height))

;; print extra stuff to the screen while solving
(define verbose #f)

;; (build-list 4 (lambda (x) (add1 x))) = '(1 2 3 4)
;; equivalent to (map fn (list 0 .. n - 1))

;; return a list of numbers from 1 to n * m
(define (all-numbers n m)
  (build-list (* n m) (lambda (x) (add1 x))))

;; return the value of a cell within the board at x,y 
(define (get-cell board x y)
  (let ((cells (board-cells board)))
    (list-ref cells (+ x (* (board-height board) (board-width board) y)))))

;; return the value of a cell within a square at x, y
(define (get-square-cell square x y)
  (let ((cells (board-cells square)))
    (list-ref cells (+ x (* (board-width square) y)))))

;; return a new board with the cell at position x,y( reversed here )
(define (update-board board y x cell)
  (when (null? cell)
    (error (format "updating board at ~a ~a ~a with empty cell\n" y x (get-cell board y x))))
  (let loop ((height 0)
	     (width 0)
	     (cells '()))
    (cond
      ((>= height (* (board-width board) (board-height board)))
       (make-board (reverse cells)
		   (board-width board)
		   (board-height board)))
      ((>= width (* (board-height board) (board-width board)))
       (loop (add1 height) 0 cells))
      ((and (= height x)
	    (= width y))
       (loop height (add1 width) (cons cell cells)))
      (else
	(loop height (add1 width) (cons (get-cell board width height) cells))))))

;; create a board and set all values to be unknown
(define (setup-board m n)
  (make-board (build-list (* m n m n) (lambda (x) (all-numbers m n)))
	      m n))

;; "string times" - return 'num' copies of 'str' appended together
;; (string* "xa" 5) = "xaxaxaxaxa"
(define (string* str num)
  (apply string-append
	 (build-list num (lambda (x) str))))

;; print a board
(provide print-board)
(define (print-board board)
  (let ((len (* (board-width board) (board-height board))))
    (let loop ((height 0)
	       (width 0))
      (cond
	((>= height len) (void))
	((>= width (sub1 len))
	 (begin
	   (printf "~a\n" (get-cell board width height))
	   (when (= (modulo (add1 height) (board-height board)) 0)
	     (printf "~a\n" (string* "-" (* len len))))
	   (loop (add1 height) 0)))
	(else
	  (begin
	    (printf "~a" (get-cell board width height))
	    (when (= (modulo (add1 width) (board-width board)) 0)
	      (printf " | "))
	    (loop height (add1 width))))))))

;; set all cells to have the value 'cell'
(define (update-all board cell)
  (let ((len (* (board-width board) (board-height board))))
    (let loop ((height 0)
	       (width 0)
	       (board board))
      (cond
	((>= height len) board)
	((>= width len)
	 (loop (add1 height) 0 board))
	(else
	  (loop height (add1 width) (update-board board width height cell)))))))

;; print square to screen
(define (print-square square)
  (let loop ((height 0)
	     (width 0))
    (cond
      ((>= height (board-height square)) (void))
      ((>= width (sub1 (board-width square)))
       (begin
	 (printf "~a\n" (get-square-cell square width height))
	 (loop (add1 height) 0)))
      (else
	(printf "~a" (get-square-cell square width height))
	(loop height (add1 width))))))

;; return the inner square within the entire board
;; on a sudoku puzzle with 3x3 inner squares and thus
;; a 9x9 board, (get-square board 0 0) is the upper left square
;; and (get-square board 2 2) is the lower right square
(define (get-square board x y)
  (let ((min-x (* x (board-width board)))
	(min-y (* y (board-height board)))
	(max-x (sub1 (+ (board-width board) (* x (board-width board)))))
	(max-y (sub1 (+ (board-height board) (* y (board-height board))))))
    (let loop ((height min-y)
	       (width min-x)
	       (cells '()))
      ;; (printf "Get square ~a ~a now ~a ~a\n" x y width height)
      (cond
	((> height max-y) (make-board (reverse cells)
				       (board-width board)
				       (board-height board)))
	((> width max-x)
	 (loop (add1 height) min-x cells))
	(else
	  (loop height (add1 width) (cons (get-cell board width height) cells)))))))

;; return a new board with the cell at x,y within square at
;; square-x, square-y equal to 'cell'
(define (update-square-cell board square-x square-y x y cell)
  (update-board board
		(+ x (* square-x (board-width board)))
		(+ y (* square-y (board-height board)))
		cell))

;; replace an entire square on the board
(define (update-square board square-x square-y square)
  (let loop ((y 0)
	     (x 0)
	     (board board))
    (cond
      ((>= y (board-height square)) board)
      ((>= x (board-width square))
       (loop (add1 y) 0 board))
      (else
	(loop y (add1 x) (update-square-cell board square-x square-y
					     x y
					     (get-square-cell square x y)))))))

;; #t if board is a solved sudoku puzzle
(provide solved?)
(define (solved? board)
  (define (check-board)
    (let ((len (* (board-width board) (board-height board))))
      (let loop ((height 0)
		 (width 0))
	(cond
	  ((>= height len) #t)
	  ((>= width len)
	   (loop (add1 height) 0))
	  (else
	    (if (> (length (get-cell board width height)) 1)
	      #f
	      (loop height (add1 width))))))))
  (define (check-squares)
    (define (check-square square)
      (let ((square-num (apply + (apply append (board-cells square))))
	    (total-num (/ (let ((n (add1 (* (board-width board)
					    (board-height board)))))
			    (* n (sub1 n)))
			  2)))
	(= square-num total-num)))
    (let loop ((height 0)
	       (width 0))
      (cond
	((>= height (board-width board)) #t)
	((>= width (board-height board))
	 (loop (add1 height) 0))
	(else
	  (and (check-square (get-square board width height))
	       (loop height (add1 width)))))))
  (and (check-board)
       (check-squares)))

;; #t if n contains any duplicate values
(define (duplicates? n)
  (let loop ((ns n))
    (if (null? ns)
      #f
      (let ((a (car ns))
	    (rest (cdr ns)))
	(if (memq a rest)
	  #t
	  (loop rest))))))

;; #t if board is not valid, i.e. same two numbers in the same column/row
(define (invalid? board)
  (define (flatten cells)
    (apply append (filter (lambda (x) (= (length x) 1)) cells)))
  (define (check-squares)
    (let loop ((height 0)
	       (width 0))
      (cond
	((>= height (board-width board)) #f)
	((>= width (board-height board))
	 (loop (add1 height) 0))
	(else
	  (let ((single-nums (flatten (board-cells (get-square board width height)))))
	    (or (duplicates? single-nums)
		(loop height (add1 width))))))))
  (define (check-board)
    (define (check-rows)
      (let loop ((height 0))
	(if (>= height (board-height board))
	  #f
	  (let ((all (flatten (build-list (board-width board)
					  (lambda (x)
					    (get-cell board height x))))))
	    (or (duplicates? all)
		(loop (add1 height)))))))
    (define (check-columns)
      (let loop ((width 0))
	(if (>= width (board-width board))
	  #f
	  (let ((all (flatten (build-list (board-height board)
					  (lambda (y)
					    (get-cell board width y))))))
	    (or (duplicates? all)
		(loop (add1 width)))))))
    (and (check-rows)
	 (check-columns)))
  (or (check-squares)
      (check-board)))

(define (rotate-list some-list n)
  (let loop ((lst some-list)
	     (count n))
    (if (= 0 count)
      lst
      (loop (append (cdr lst) (list (car lst))) (sub1 count)))))

;; turn '(1 2 3 4) n = 2 into '((1 2) (3 4))
(define (partition-list some-list n)
  (let loop ((so-far '())
	     (all '())
	     (rest some-list)
	     (count n))
    (cond
      ((null? rest) (reverse (cons (reverse so-far) all)))
      ((= count 0)
       (loop '() (cons (reverse so-far) all) rest n))
      (else
	(loop (cons (car rest) so-far) all (cdr rest) (sub1 count))))))
#;
(printf "1234 = ~a\n" (rotate-list '(1 2 3 4) 1))

#;
(printf "12345678 = ~a\n" (partition-list '(1 2 3 4 5 6 7 8) 2))

;; rotate a sudoku square by x,y
;; a square like
;; 1 2
;; 3 4
;; rotated by 1 x value becomes
;; 2 1
;; 4 3
;; and the original square rotated by 1 y value becomes
;; 3 4
;; 1 2
(define (rotate-square cells width height x-rotate y-rotate)
  (map (lambda (x) (list x))
       (apply append
	      (rotate-list
		(map (lambda (l) (rotate-list l x-rotate))
		     (partition-list cells width))
		y-rotate))))

#;
(printf "~a\n" (rotate-square '((1) (2) (3) (4)) 2 2 1 0))

;; randomize a list using knuth's subtractive method
(define (randomize lst)
  (let ((v (list->vector lst)))
    (let loop ((max (sub1 (vector-length v))))
      (if (= 0 max)
	(vector->list v)
	(begin
	  (let ((place (random max)))
	    (let ((tmp (vector-ref v place)))
	      (vector-set! v place (vector-ref v max))
	      (vector-set! v max tmp)))
	  (loop (sub1 max)))))))

#;
(printf "~a\n" (randomize '(1 2 3 4 5)))

;; create an already solved sudoku puzzle
(provide create-sudoku)
(define (create-sudoku width height)
  (define (fill-in board)
    (define randomized-square (randomize (all-numbers width height)))
    (let loop ((y 0)
	       (x 0)
	       (board board))
      (cond
	((>= y width) board)
	((>= x height)
	 (loop (add1 y) 0 board))
	(else
	  (loop y (add1 x) (update-square board x y (make-board (rotate-square
								  randomized-square
								  width height y x)
								width height)))))))

  (let ((board (fill-in (setup-board width height))))
    board))

;; make some cells in the board unknown
(define (toughen board num)
  (let loop ((num num)
	     (board board))
    (if (= num 0)
      board
      (let ((x (random (* (board-width board) (board-height board))))
	    (y (random (* (board-width board) (board-height board)))))
	(loop (sub1 num)
	      (update-board board x y
			    (all-numbers (board-width board) (board-height board))))))))

(define (toughen-50% board)
  (toughen board (floor (/ (* (board-width board)
			      (board-height board))
			   1))))

;; return a sudoku puzzle that should be solvable
#;
(define (create-sudoku width height)
  (define (make)
    (let ((board (setup-board width height))
	  (sets (add1 (random (* width height))))
	  ;; (sets (* width height width height width height))
	  )
      (define (random-width) (random width))
      (define (random-height) (random height))
      (define (random-number) (add1 (random (* width height))))
      (let loop ((square-x (random-width))
		 (square-y (random-height))
		 (x (random-width))
		 (y (random-height))
		 (num sets)
		 (board board))
	(if (= num 0)
	  board
	  (loop (random-width) (random-height) (random-width) (random-height)
		(sub1 num)
		(update-square-cell board square-x square-y x y (list (random-number))))))))
  (let loop ((c (make)))
    (if (invalid? c)
      (loop (make))
      c)))

;; only elements from lst2 that don't appear in lst1
(define (subtract-list lst1 lst2)
  (let loop ((ns lst2)
	     (ok '()))
    (if (null? ns)
      ok
      (let ((a (car ns))
	    (rest (cdr ns)))
	(if (memq a lst1)
	  (loop rest ok)
	  (loop rest (cons a ok)))))))

;; go through each cell and update it with all the values that
;; it cannot possibly contain. This consists of finding invalid
;; values in each row, column and square
(define (update-cells board)
  ;; return all the used values in a column
  (define (invalid-rows board x y)
    (apply append
	   (build-list (* (board-width board) (board-height board))
		       (lambda (height)
			 (if (= height y)
			   '()
			   (let ((cell (get-cell board x height)))
			     (if (null? (cdr cell))
			       cell
			       '())))))))
  ;; return all the used values in a row
  (define (invalid-columns board x y)
    (apply append
	   (build-list (* (board-height board) (board-width board))
		       (lambda (width)
			 (if (= width x)
			   '()
			   (let ((cell (get-cell board width y)))
			     (if (null? (cdr cell))
			       cell
			       '())))))))
  ;; return all the used values in a square
  (define (invalid-square board x y)
    (let ((square (get-square board
			      (round (floor (/ x (board-width board))))
			      (round (floor (/ y (board-height board)))))))
      (apply append (filter (lambda (x) (= (length x) 1)) (board-cells square)))))
  (define (update-cell board x y)
    (update-board board x y
		  (if (= 1 (length (get-cell board x y)))
		    (get-cell board x y)
		    (let ((choices (append (invalid-rows board x y)
					   (invalid-columns board x y)
					   (invalid-square board x y))))
		      #;
		      (printf "Invalid rows for ~a ~a ~a = ~a\n" x y (invalid-rows board x y)
			      (sort
				(subtract-list choices
					       (build-list (* (board-width board)
							      (board-height board))
							   (lambda (x) (add1 x))))
				<))
		      ;; values can be whatever is not in 'choices'
		      ;; sort it to make it look pretty
		      (sort
			(subtract-list choices
				       (all-numbers (board-width board)
						    (board-height board)))
			<)))))
  (let ((len (* (board-width board) (board-height board))))
    (let loop ((height 0)
	       (width 0)
	       (board board))
      (cond
	((>= height len) board)
	((>= width len)
	 (loop (add1 height) 0 board))
	(else
	  (loop height (add1 width) (update-cell board width height)))))))

;; create a n X m sudoku board from an input port. port should contain n*n*m*m characters.
;; characters '1'-'9' are values on the board and 'X' is the unknown value.
;; n*m cannot be larger than 9
(define (read-sudoku width height port)
  (let loop ((left (* width height width height))
	     (cells '()))
    (if (= left 0)
      (make-board (reverse cells) width height)
      (let ((c (read-char port)))
	;; (printf "Read ~a\n" c)
	(cond
	  ((eof-object? c) (error "not enough characters in sudoku stream"))
	  ((char=? c #\X) (loop (sub1 left) (cons (all-numbers width height) cells)))
	  (else (loop (sub1 left) (cons (list (- (char->integer c) (char->integer #\0)))
					cells))))))))

(define (read-3x3-sudoku port)
  (read-sudoku 3 3 port))
(define (read-2x3-sudoku port)
  (read-sudoku 2 3 port))

(define (read-input-sudoku port)
  (define (get regexp)
    (let ((match (pregexp-match regexp port)))
      (if (not match)
	(error (format "could not match ~a\n" regexp))
	(bytes->string/utf-8 (car match)))))
  (let* ((width (string->number (get "^\\d+")))
	 (_ (get "\\s*"))
	 (height (string->number (get "^\\d+"))))
    ;; (printf "width is ~a height ~a\n" width height)
    (let loop ((n (* width height width height))
	       (cells '()))
      (if (= n 0)
	(make-board (reverse cells) width height)
	(let ((next (get "\\d+|_"))
	      (_ (get "\\s*\n*")))
	  (cond
	    ((string=? "_" next)
	     (loop (sub1 n) (cons (all-numbers width height)
				  cells)))
	    (else
	      (loop (sub1 n) (cons (list (string->number next)) cells)))))))))

;; perform 1 solution step
(define (solve1 board)
  (update-cells board))

;; #t if b1 is the same as b2
(define (board=? b1 b2)
  (equal? (board-cells b1) (board-cells b2)))

(define (board->list board)
  (apply append (board-cells board)))

;; pick a cell that has multiple choices of values and then choose a value
(define (choose-cell board)
  (let loop ((cells '())
	     (rest (board-cells board)))
    (cond
      ((null? rest) (error "No cells to choose from"))
      (else
	(let ((a (car rest))
	      (b (cdr rest)))
	  ;; if cell isn't solved
	  (if (not (= (length a) 1))
	    (make-board (append (reverse (cons (list (car a)) cells)) b)
			(board-width board)
			(board-height board))
	    (loop (cons a cells) b)))))))

(define-struct cell (x y num))
(define (unknown-cells board)
  (define (more-cells cells x y)
    (let ((cell (get-cell board x y)))
      (if (= (length cell) 1)
	cells
	(append cells
		(map (lambda (num) (make-cell x y num)) cell)))))
  (let ((len (* (board-width board) (board-height board))))
    (let loop ((cells '())
	       (x 0)
	       (y 0))
      (cond
	((>= y len) cells)
	((>= x len)
	 (loop cells 0 (add1 y)))
	(else
	  (loop (more-cells cells x y) (add1 x) y))))))

;; solve a sudoku puzzle. raises an error if anything goes wrong
;; at each step the solver will try to reduce the board by
;; narrowing the list of possible values a cell could have given
;; the constraints of its rows/columns/square.
;; if a board cannot be reduced then a cell is chosen and automatically
;; reduced to a single number
(provide solve)
(define (solve board)
  (let loop ((board board))
    (when verbose
      (printf "Board is now\n")
      (print-board board)
      (newline))
    (cond
      ((solved? board) board)
      (else
	(let ((next-board (solve1 board)))
	  (if (board=? next-board board)
	    (let ((all-possible (unknown-cells board)))
	      (ormap (lambda (cell)
		       (let ((x (cell-x cell))
			     (y (cell-y cell))
			     (num (cell-num cell)))
			 ;; (printf "Trying board with ~a,~a = ~a\n" x y num)
			 (with-handlers
			    (((lambda (x) (exn:fail? x))
			      (lambda (x) #f)))
			    (loop (update-board board x y (list num))))))
		     all-possible))
	    ;; (loop (choose-cell board))
	    (loop next-board)))))))

#;
(let ((board (setup-board 2 3)))
  ;; (printf "~a\n" (board-cells board))
  ;; (newline)
  (print-board board)
  (printf "Solved? ~a\n" (solved? board))
  (let ((b1 (update-all (setup-board 2 2) '(1))))
    (print-board b1)
    (printf "Solved? ~a\n" (solved? b1)))
  (print-square (get-square board 0 0))
  (print-square (get-square (update-board board 0 1 '(1)) 0 0))
  (print-square (get-square (update-square-cell board 0 0 1 1 '(1)) 0 0))
  (printf "Solved? ~a\n" (solved? (update-square-cell (update-all board '(2)) 0 0 0 0 '(1))))
  ;; (get-square board 1 0)
  ;; (get-square board 0 1)

  #;(printf "~a\n" (board-cells (update-board board 0 0 '(1))))
  (newline)
  #;(print-board (update-board board 1 2 '(1))))

  #;
(let ((puzzle (create-soduku 2 2)))
  (print-board puzzle)
  (printf "Solve 1\n")
  (print-board (solve puzzle))
  (printf "Solve 2\n")
  (print-board (solve (solve puzzle)))
  )

#;
(let* ((puzzle (setup-board 2 2))
       (puzzle (update-square-cell puzzle 0 0 0 1 '(3)))
       (puzzle (update-square-cell puzzle 1 0 0 1 '(2)))
       (puzzle (update-square-cell puzzle 1 0 0 0 '(3)))
       )
  (print-board puzzle)
  (printf "Solve 1\n")
  (print-board (solve puzzle))
  (printf "Solve 2\n")
  (print-board (solve (solve puzzle))))


#;
(print-board (solve (create-sudoku 2 2)))
#;
(let ((r (read-3x3-sudoku (open-input-file "p"))))
  (print-board r)
  (printf "~a\n" (board->list (solve r))))

(define (tests)
  (define (verify board output)
    (equal? (board->list (solve board)) output))

  (define (verify-3x3 input output)
    (verify (read-3x3-sudoku (open-input-string input)) output))

  (define (verify-2x3 input output)
    (verify (read-2x3-sudoku (open-input-string input)) output))

  (define (verify-random n m)
    (let* ((b (create-sudoku n m))
	   (hard (toughen-50% b)))
      (printf "Random ~ax~a ~a\n" n m (verify hard (board->list b)))))

  (printf "1. 3x3 ~a\n" (verify-3x3 (string-append "XXX5864XX"
						   "65XXXXX82"
						   "XX3XX4XXX"
						   "X4976XX1X"
						   "72XXXXX69"
						   "X6XX3972X"
						   "XXX2XX8XX"
						   "41XXXXX53"
						   "XX6357XXX")
				    '(1 9 2 5 8 6 4 3 7
				      6 5 4 9 7 3 1 8 2
				      8 7 3 1 2 4 6 9 5
				      3 4 9 7 6 2 5 1 8
				      7 2 8 4 1 5 3 6 9
				      5 6 1 8 3 9 7 2 4
				      9 3 5 2 4 1 8 7 6
				      4 1 7 6 9 8 2 5 3
				      2 8 6 3 5 7 9 4 1)))

  (printf "2. 3x3 ~a\n" (verify-3x3 (string-append "X82XX96XX"
						   "X6XX359XX"
						   "X93XXX7X8"
						   "35XXXX41X"
						   "2X9XXX3X7"
						   "X17XXXX95"
						   "8X1XXX23X"
						   "XX621XX7X"
						   "XX54XX18X")
				    '(1 8 2 7 4 9 6 5 3
				      7 6 4 8 3 5 9 2 1
				      5 9 3 1 6 2 7 4 8
				      3 5 8 6 9 7 4 1 2
				      2 4 9 5 8 1 3 6 7
				      6 1 7 3 2 4 8 9 5
				      8 7 1 9 5 6 2 3 4
				      4 3 6 2 1 8 5 7 9
				      9 2 5 4 7 3 1 8 6)))

  (printf "3. 2x3 ~a\n" (verify-2x3 (string-append
				      "1X453X"
				      "3X16X5"
				      "XX3X1X"
				      "X1XX63"
				      "5X6X4X"
				      "X62X5X")
				    '(1 2 4 5 3 6
				      3 4 1 6 2 5
				      6 5 3 2 1 4
				      2 1 5 4 6 3
				      5 3 6 1 4 2
				      4 6 2 3 5 1)))

  #;
  (let ((s (create-sudoku 3 3)))
    (print-board s)
    (newline)
    (print-board (toughen-50% s))
    )



  (verify-random 1 1)
  (verify-random 1 2)
  (verify-random 1 3)
  (verify-random 1 4)

  (verify-random 2 1)
  (verify-random 2 2)
  (verify-random 2 3)
  (verify-random 2 4)

  (verify-random 3 1)
  (verify-random 3 2)
  (verify-random 3 3)
  (verify-random 3 4)

  (verify-random 4 1)
  (verify-random 4 2)
  (verify-random 4 3)
  (verify-random 4 4)

  #;
  (let loop ((n 1))
    (when (< n 10)
      (printf "~a\n" n)
      (time (solve (toughen-50% (toughen-50% (create-sudoku n n)))))
      (loop (add1 n))))

  )

;; (print-board (solve (create-sudoku 2 3)))

;; (print-board (create-sudoku 3 3))
(let ((solved (solve (read-input-sudoku (current-input-port)))))
  (if solved
    (print-board solved)
    (printf "No solution\n")))