(require (lib "etc.ss")) ; For build-list
(require (lib "list.ss")) ; For filter

;; ---- RANDOM LIST OPS ----
(define replace
  (lambda (lst n val) (map (lambda (i v) (if (= i n) val v)) (build-list (length lst) (lambda (x) x)) lst)))

(define take
  (lambda (lst n) (reverse (list-tail (reverse lst) (- (length lst) n)))))

;; ---- TREE STRUCTURES ----
; Treat the list as a stack
(define stack-next
  (lambda (lst) (values (car lst) (cdr lst))))

; Treat the list as a queue
(define queue-next
  (lambda (lst) (values (list-ref lst (- (length lst) 1)) (reverse (cdr (reverse lst))))))

; All successor move-lists for a given move-list
(define successor-moves
  (lambda (moves) (map (lambda (x) (cons x moves))
                       '(
                         (fill 0)
                         (fill 1)
                         (fill 2)
                         (empty 0)
                         (empty 1)
                         (empty 2)
                         (pour 0 1)
                         (pour 0 2)
                         (pour 1 0)
                         (pour 1 2)
                         (pour 2 0)
                         (pour 2 1)
                         ))))


;; Calculate the results of a pour
(define pour-helper
  (lambda (state max from to) (let* ((tomax (list-ref max to))
                                     (fromstate (list-ref state from))
                                     (tostate (list-ref state to))
                                     (toleft (- tomax tostate))
                                     (stage1 (replace state from (if (>= toleft fromstate) 0 (- fromstate toleft)))) ; TODO: use min/max for these
                                     (stage2 (replace stage1 to (if (>= toleft fromstate) (+ tostate fromstate) tomax)))
                                     )
                                stage2)))
                                

(define print-list
  (lambda (lst lbl) (if (null? lst) 
                        (printf "~a: '()~n" lbl) 
                        (begin 
                          (printf "~a: (~a" lbl (car lst)) 
                          (for-each (lambda (x) (printf "~n~a   ~a" (make-string (string-length lbl) #\space) x)) (cdr lst))
                          (printf ")~n")
                          )
                        )))

;; Compile a set of moves into a state
(define compile-moves
  (let ((memo (make-hash-table 'equal)))
    (lambda (moves max cur) (let ((hashed (hash-table-get memo (list moves max) #f)))
                              (cond
                                ((null? moves) cur)
                                (hashed hashed)
                                (else (let* ((state (compile-moves (cdr moves) max cur))
                                             (move (car moves))
                                             (rv (cond
                                                   ((equal? move '(fill 0)) (list (car max) (cadr state) (caddr state))) ; TODO: Rework using replace
                                                   ((equal? move '(fill 1)) (list (car state) (cadr max) (caddr state)))
                                                   ((equal? move '(fill 2)) (list (car state) (cadr state) (caddr max)))
                                                   ((equal? move '(empty 0)) (list 0 (cadr state) (caddr state)))
                                                   ((equal? move '(empty 1)) (list (car state) 0 (caddr state)))
                                                   ((equal? move '(empty 2)) (list (car state) (cadr state) 0))
                                                   ((equal? (car move) 'pour) (pour-helper state max (cadr move) (caddr move)))
                                                   (else (raise 'invalid-move))
                                                   ))
                                             )
                                        (hash-table-put! memo (list moves max) rv)
                                        rv)))))))
      
; Main workhorse function
; Here be Dragons.
(define search
  (let ((memo (make-hash-table 'equal)))
    (lambda (moveslst max target d nextfn)
      (let-values (((moves rest) (nextfn moveslst))) ;(printf "State: ~a~n" (reverse moves))
        (let ((state (compile-moves moves max '(0 0 0))) (nmoves (length moves)))
          (if (equal? state target) 
              (raise (list #t (reverse moves) memo)) ; We did it!
              (begin
                (hash-table-put! memo state nmoves)
                (let* ((suc (successor-moves moves)) ; All successor move lists
                       (good-suc (filter (lambda (x) (and (let ((h (hash-table-get memo (compile-moves x max '(0 0 0)) #f))) 
                                                            (if h (> h (length x)) #t)) ; If ths state is in the memo, only pass it if this move-list is shorter
                                                          (<= (length x) d))) ; This is the depth control
                                         suc)) ; Only legal move lists
                       (nextmoveslst (append good-suc rest))
                       ) ;(print-list (map reverse suc) "Suc")(print-list (map reverse good-suc) "Good Suc")
                  (if (null? nextmoveslst)
                      (raise (list #f '() memo)) ; Bail out
                      (search nextmoveslst max target d nextfn) ; Yay, tail recursion!
                      )
                  )
                )
              )
          )
        )
      )
    )
  )

; Same as above sans memorizing (not used)
(define slow-search
  (lambda (moveslst max target d nextfn) 
    (let-values (((moves rest) (nextfn moveslst))) ;(printf "State: ~a~n" (reverse moves))
      (let ((state (compile-moves moves max '(0 0 0))))
        (if (equal? state target) 
            (raise (list #t (reverse moves))) ; We did it!
            (let* ((suc (successor-moves moves)) ; All successor move lists
                   (good-suc (filter (lambda (x) (<= (length x) d)) suc)) ; Only legal move lists
                   (nextmoveslst (append good-suc rest))
                   ) ;(print-list (map reverse suc) "Suc")(print-list (map reverse good-suc) "Good Suc")
              (if (null? nextmoveslst)
                  (raise (list #f '())) ; Bail out
                  (slow-search nextmoveslst max target d nextfn)
                  )
              )
            )
        )
      )
    )
  )

; Wrapper to intercept the return value from search
(define do-search
  (lambda (max target d nextfn)
    (with-handlers ((list? (lambda (exn) 
                             (hash-table-for-each (caddr exn) (lambda (k v) (hash-table-remove! (caddr exn) k))) 
                             (if (car exn) (cadr exn) #f))
                           ))
      (search (list '()) max target d nextfn)
      "ERROR"))) ; We should never see this 

;; ---- TOP LEVEL SEARCHES ----
; Breadth first search
(define bfs-search
  (lambda (max target d) (do-search max target d queue-next)))

; Depth first search
(define dfs-search
  (lambda (max target d) (do-search max target d stack-next)))
  
; Iterated DFS
(define iter-dfs-search
  (lambda (max target d) (letrec ((helper (lambda (i) (if (> i d) 
                                                          #f 
                                                          (let ((rv (dfs-search max target i)))
                                                            (if rv 
                                                                rv 
                                                                (helper (+ i 1))
                                                                )
                                                            )
                                                          )
                                            )))
                           (helper 1))))

;; ---- RENDER THE OUTPUT ----
; Convert a list of moves (of the form (action x [y])) to a list of strings
(define format-moves
  (lambda (moves) (if moves
                      (map (lambda (m i) (apply format (string-append "~a. " (cond
                                                                                ((eq? (car m) 'fill) "Fill jug ~a.")
                                                                                ((eq? (car m) 'empty) "Emptry jug ~a.")
                                                                                ((eq? (car m) 'pour) "Pour jug ~a into jug ~a.")
                                                                                ))
                                                (cons i (cdr m))))
                                moves
                                (build-list (length moves) (lambda (i) (+ i 1))))
                      (list "No solution found"))))

; Print a list of strings to the screen
(define print-moves
  (lambda (moves) (for-each (lambda (s) (display s)(newline)) moves)))

;; ---- THE FUNCTIONS ---
(define dfswj
  (lambda args (format-moves (apply iter-dfs-search args))))

(define bfswj
  (lambda args (format-moves (apply bfs-search args))))
 
;; ---- TEST STUFF ----
;(define x (dfs-search '(1 2 3) '(1 1 1) 3))
;(for-each (lambda (s) (display s)(newline)) (format-moves x))
;(define x (dfswj '(3 4 10) '(0 2 0) 10))
;
;(dfswj '(3 4 10) '(3 4 10) 10)
;(bfswj '(3 4 10) '(3 4 10) 10)
;(dfswj '(3 4 10) '(2 2 2) 10)
;(bfswj '(3 4 10) '(2 2 2) 10)
;(dfswj '(3 4 10) '(1 2 0) 10)
;(bfswj '(3 4 10) '(1 2 0) 10)
;(dfswj '(1 2 3) '(1 1 1) 10)
;(bfswj '(1 2 3) '(1 1 1) 10)