;; astar algorithm for 8-puzzle
;;
;; basic data structure to hold representation of a state:
;;   list indicating which piece is in which slot in the puzzle.
;;   puzzle coordinates are:
;;   
;;    0    1   2
;;    3    4   5
;;    6    7   8
;;
(load-option 'hash-table)	; MIT scheme support for hash tables
(load-option 'wt-tree)		; MIT scheme support for weight-balanced trees


(define astar-debug #t)           ; global debugging flag
				  ; prints messages every 100 nodes

(define (astar start-state goal-state get-children heuristic)
    (let ((open (make-astar-list #t))     ; open list with p.q.
	  (closed (make-astar-list #f))   ; closed list - no p.q.
	  (newboards ())                  ; holds successors
	  (currentboard()))               ; current state

      (set! heuristic-cnt 0)              ; initialize to 0

      ; add the start state to the open list
      (open 'add 
	    (8puzzle-record 
	     start-state 
	     '() 
	     0 
	     heuristic 
	     goal-state))

      ; named let used for recursion 

      (let astar-search ((step 1))
	; some debugging stuff
	(if astar-debug
	    (astar-print-status step open closed))

	; get the best state from the open list
	(set! currentboard (open 'getnext))

	; get a list of successor states
	(set! newboards (get-children (cadr currentboard)))

	; check for the goal state
	(if (equal? (cadr currentboard) goal-state)
	    ; DONE - print results
	    (print-path currentboard closed #t)
	    ; NOT DONE 
	    (begin
	      ; add current state to closed list
	      (closed 'add currentboard)

	      ; add each new board to queue if not in open or closed.
	      (map (lambda (b)
		     (let (
			   (onclosed (closed 'lookup b))
			   (onopen   (open   'lookup b)))
		       (cond
			(onclosed
			 ; state is on the closed list
			 ; remove the entry from closed and
			 ; put the better one back on closed
			 (closed 'remove b)
			 (if (< (cadddr onclosed)
				(+ 1 (cadddr currentboard)))
			     ; old one is better
			     (closed 'add onclosed)
			     ; new one is better
			     (closed 'add 
				     (8puzzle-record 
				      b (cadr currentboard)
				      (+ 1 (cadddr currentboard))
				      heuristic goal-state))))

			(onopen
			 ; state is on the open list
			 ; remove the entry from open and
			 ; put the better one back on open
			 (open 'remove b)
			 (if (< (cadddr onopen)
				(+ 1 (cadddr currentboard)))
			     ; old one is better
			     (open 'add onopen)
			     ; new one is better
			     (open 'add 
				   (8puzzle-record 
				    b (cadr currentboard)
				    (+ 1 (cadddr currentboard))
				    heuristic goal-state))))

			(else
			 ; state wasn't on open or closed
			 ; add it to the open list
			 (open 'add 
			       (8puzzle-record 
				b (cadr currentboard)
				(+ 1 (cadddr currentboard))
				heuristic goal-state))))))
		   newboards)
	      (astar-search (+ step 1)))))))


; Print out the status of the run every 100 steps

(define astar-print-status
  (lambda (step open closed)
    (if (= 0 (modulo step 100))
	(begin
	  (display "STEP #")
	  (display step)
	  (display ": OPEN ")
	  (display (open 'length))
	  (display " CLOSED ")
	  (display (closed 'length))
	  (display "\n")))))

; ======================================================
; scheme support for A* lists using hash table and weight balanced trees.
;
; this can all be done with simple lists, but the time per operation
; is way too slow to be useful for large lists
;
; A* lists are composed of keyed records, and can also
; be used as a priority queue. A hash table is used to
; provide constant time lookups, and a weight-balanced 
; tree is used for a fast priority queue. The full record
; for each state is stored in the hash table, only the
; value of the f-function and the state representation itself
; are stored in the priority queue.
;
; the structure of each state record is a list of 4 elements,
;  1. number that indicates the value of f for this state.
;  2. representation of the state
;  3. representation of the parent state (predecesor)
;  4. number indicating the path length from the start state
;        to the state.
;
; For the 8-puzzle, each state is represented by a list containing
; 9 elements, each an integer. The first element of the list indicates the
; tile that is currently in the upper left corner of the puzzle. The second
; is the middle tile on the top row, etc.
; the puzzle state:
;
;  5 2 3
;  0 4 7
;  1 6 8
;
; is represented by the list '(5 2 3 0 4 7 1 6 8)
;
; 




; -----------------------------------------------
; priority queue support
;

; add a new record to the priority queue

(define astar-pq-add
  (lambda (wt rec)
    (wt-tree/add! wt rec rec)))

; remove a record from the priority queue

(define astar-pq-remove
  (lambda (wt rec)
    (wt-tree/delete! wt rec)))

; get next key from the priority queue

(define astar-pq-getnext
  (lambda (wt)
    (wt-tree/min wt)))

; key comparison proc for wt tree and 8 puzzle
; each key is the cons of the f value and the puzzle state
;
(define wtcomp
  (lambda (a b)
    (cond
     ((null? a) #f)
     ((equal? (car a) (car b))
      (wtcomp (cdr a) (cdr b)))
     ((equal? (car a) 'space) #t)
     ((equal? (car b) 'space) #f)
     (else (< (car a) (car b))))))
    



;============================================================
; support for list used by A* algorithm for open and closed lists
;
; make-astar-list takes a single parameter logical that indicates
; whether the list will be used as a priority queue.
; 
; For the A* algorithm, we need to extract states from the open list
; by f-value (we always look at the state in the open list with the
; best (smallest) f-value.
;
; The closed list does not need the priority queue (we only use it to
; do lookups), so we don't want to waste time maintaining the queue
; if we don't need it.
;
; operations on the astar list:
; add - adds a new record: (f-val state parent pathlength)
; lookup - returns a record given a state (or ())
; getnext - returns next best record and deletes it from the list
; remove - deletes a record from the list given a state
; print - used for debugging
; clear - cleans out everything (emptys the list)
; length - returns the number of elements on the list
;
(define make-astar-list
  (lambda (use-priority-queue)
    (let ((ht (make-equal-hash-table))        ; the hash table
	  (use-pq use-priority-queue)         ; the priority queue flag
	  (pq (make-wt-tree                   ; the priority queue
	       (make-wt-tree-type wtcomp)))
          (b ()))                             ; temp variable
      (lambda (op . args)

	(cond

	 ((eqv? op 'add)                       ; adding a new record
	  (set! b (cadar args))                ; b is the state
	  (if (hash-table/get ht b #f)         ; check for duplicates
	      (begin
		(display "DUPLICATE ADD ")
		(display b)
		(display "\n")))

	  (hash-table/put! ht b (car args))    ; add to the hash table
	  ; if we are maintaining the priority queue - add it there as well
	  (if use-pq (astar-pq-add pq (cons (caar args) b))))



	 ((eqv? op 'lookup)                     ; lookup record given a state
	  (hash-table/get ht  (car args) #f))


	 ((eqv? op 'getnext)                    ; get next record from p.q.
	  (if use-pq
	      (let ((x (astar-pq-getnext pq)))  ; get next state
		(astar-pq-remove pq x)          ; remove from the p.q.
		(set! b (hash-table/get ht (cdr x) #f))   ; get whole record
		(hash-table/remove! ht (cdr x))           ; remove from h.t.
		b)
	      (begin
		(display "INVALID OPERATION - NO PQ\n")
		(cdr ()))))

	 ((eqv? op 'remove)                     ; remove record given a state
	  (set! b (hash-table/get ht (car args) #f)) ; b is record
	  (hash-table/remove! ht  (car args))        ; remove from h.t.
	  (if use-pq                                 ; and from p.q.
	      (astar-pq-remove pq (cons (car b) (cadr b)))))


	 ((eqv? op 'print)
	  (if use-pq
	      (let ((c 1))
		(wt-tree/for-each 
		 (lambda (key value)
		   (display c)
		   (display ": ")
		   (display key)
		   (display " | " )
		   (display (hash-table/get ht (cdr key) #f))
		   (display "\n")
		   (set! c (+ c 1)))
		 pq))
	      (begin
		(display "INVALID OPERATION - NO PQ\n")
		(cdr ()))))
	 


	 ((eqv? op 'clear)
	  (if use-pq                         ; clean out the p.q.
	      (wt-tree/for-each 
	       (lambda (key value)
		 (wt-tree/delete! pq key))
	       pq))
	  (hash-table/clear! ht))            ; and the has table


	 ((eqv? op 'length)                   
	  (hash-table/count ht))             ; return the number of records

	 (else 
	  (display "ASTAR-LIST ")
	  (display op)
	  (display "\n")))))))

; ===============================================================
;
; I/O support for listing a solution (path)

(define 8puzzle-drawrow
  (lambda (board coords)
    (if (null? coords) ()
	(let ((char (list-ref board (car coords))))
	  (display  (if (equal? char 'space)
			" "
			char))
	  (display " ")
	  (8puzzle-drawrow board (cdr coords))))))
		   
;;; draw a puzzle state (board)

(define 8puzzle-drawboard
  (lambda (board)
    (8puzzle-drawrow board '(0 1 2))
    (display "\n")
    (8puzzle-drawrow board '(3 4 5))
    (display "\n")
    (8puzzle-drawrow board '(6 7 8))
    (display "\n\n")))

; create a path from the board to the start state
; the path is given as a list of states

(define create-path
  (lambda (board closed)
    (if (null? board) ()
	(let ((parent (closed 'lookup (caddr board))))
	      (if parent
		  (append  (create-path parent closed) (list (cadr board)))
		  (list (cadr board)))))))
	
; print out results of a run
; uncomment to see the solution path

(define print-path
  (lambda (brec closed flag)
    (if flag (drawsolution (create-path brec closed)))
    (display "; ")
    (display (- (length (create-path brec closed)) 1))
    (display " MOVES   ")
    (display heuristic-cnt)
    (display " STATES EVALUATED\n")))

;
; draws the puzzle states given a path (list of states)

(define drawsolution
  (lambda (boards)
    (if (not (null? boards))
	(begin
	  (8puzzle-drawboard (car boards))
	  (drawsolution (cdr boards))))))


; construct and 8 puzzle record given a state, parent state,
; pathlength to the state, heuristic function and goal state
(define 8puzzle-record
  (lambda (board parent pathlen heuristic goal)
    (set! heuristic-cnt (+ 1 heuristic-cnt))
    (list (+ pathlen (heuristic board goal)) board parent pathlen)))

; variable used to keep track of the number of states examined
; during a run.

(define heuristic-cnt 0)