;;;; -*- mode:Scheme -*- ;;;;

;;; CSP solution.

;;;;;;;;;;;;;;;;;
;;; DATA STRUCTURES
;;;;;;;;;;;;;;;;;

(define-structure (csp)
  number-of-variables
  variables				; list of variable numbers
  arcs					; list of arcs, each (head tail constraint id)
  arcs-from				; vector of arcs indexed by head-var
  arcs-into				; vector of arcs indexed by tail-var
  arcs-added-for-new-var		; vector of arcs, used in checking consistency
  variable-domains			; vector of lists of var values
  variable-names			; vector of symbols = names of vars
  )

(define-structure (arc)
  tail
  head
  constraint-function
  index
  )

;;; Global variable that hold the current constraint graph.
(define *csp* (make-csp 0 '() '()
			(make-vector 0) (make-vector 0) (make-vector 0)
			(make-vector 0) (make-vector 0)))

;;; A flag to control behavior of consistent?
(define *check-all-arcs* #f)

;;; A flag to control printing
(define *verbose* #f)
(define *verbose-revise* #f)

;;; To keep track of number of operations done.
(define *number-of-arc-tests* 0.)

;;; Accessing domains of vars
(define (domain-of var) (vector-ref (csp-variable-domains *csp*) var))
(define (set-var-domain! var value)
  (vector-set! (csp-variable-domains *csp*) var value))

;;; Accessing the names of vars
(define (name-of var) (vector-ref (csp-variable-names *csp*) var))

;;; Accessing the arcs indexed by the variables they connect.
(define (arcs-from var) (vector-ref (csp-arcs-from *csp*) var))
(define (arcs-into var) (vector-ref (csp-arcs-into *csp*) var))
(define (set-arcs-from! var arcs)
  (vector-set! (csp-arcs-from *csp*) var arcs))
(define (set-arcs-into! var arcs)
  (vector-set! (csp-arcs-into *csp*) var arcs))

;;; Accessing the arcs incrementally to avoid duplicate checking.
(define (new-arcs-for-var var) (vector-ref (csp-arcs-added-for-new-var *csp*) var))
(define (set-new-arcs-for-var! var arcs)
  (vector-set! (csp-arcs-added-for-new-var *csp*) var arcs))

;;; Operations on all the domains.

;;; Retuns a list of the current domain values, so they can be restored
;;; when backtracking.
(define (get-list-of-domains)
  (vector->list (csp-variable-domains *csp*)))

;;(define (restore-domains saved)
;; (set-csp-variable-domains! *csp* (list->vector saved)))
;; Same effect but churns less space since it does not create new vector.
(define (restore-domains saved)
  (define (restore-domains-loop s i)
    (if (>= i (csp-number-of-variables *csp*))
	#f
	(begin
	  (set-var-domain! i (first s))
	  (restore-domains-loop (rest s) (+ 1 i)))))
  (restore-domains-loop saved 0))

;;;;;;;;;;;;;;;;;
;;; GENERAL CONSTRAINT PROPAGATION CODE.
;;;;;;;;;;;;;;;;;

(define (initialize names-and-domains arcs)
  ;;
  ;; The names-and-domains is a list of lists (one for each variable).
  ;; Each list's first element is the name of the variable, the rest
  ;; of the list is the domain, that is, a list of legal values.
  ;; The arcs are lists accessed by arc-tail (first), arc-head (second),
  ;; arc-constraint-function (third) and arc-index (fourth).
  ;;
  (display* "Initializing...")
  (set-csp-arcs! *csp* arcs)
  (set-csp-number-of-variables! *csp* (length names-and-domains))
  (set-csp-variables! *csp* (make-index-list (csp-number-of-variables *csp*)))
  (set-csp-variable-names! *csp* (list->vector (map first names-and-domains)))
  (set-csp-variable-domains! *csp* (list->vector (map rest names-and-domains)))
  (set-csp-arcs-from! *csp* (make-vector (csp-number-of-variables *csp*) '()))
  (for-each
   (lambda (arc)
     (set-arcs-from! (arc-tail arc) 
		     (cons arc (arcs-from (arc-tail arc)))))
   arcs)
  ;; The arcs indexed by their head-var value is needed to implement
  ;; REVISE efficiently.  The arcs stored in i-th entry of this vector
  ;; are the arcs affected by the change in the domain of the i-th
  ;; variable. 
  (set-csp-arcs-into! *csp* (make-vector (csp-number-of-variables *csp*) '()))
  (for-each
   (lambda (arc)
     (set-arcs-into! (arc-head arc) 
		     (cons arc (arcs-into (arc-head arc)))))
   arcs)
  ;; The (csp-arcs-added-for-new-var *csp*) vector is needed to implement
  ;; CONSISTENT? efficiently. The i-th entry are the arcs that
  ;; constrain the i-th variable relative to the first i-1 variables.
  (set-csp-arcs-added-for-new-var! *csp* (make-vector (csp-number-of-variables *csp*) '()))
  (for-each
   (lambda (i)
     (for-each
      (lambda (arc)
	(if (and (<= (arc-tail arc) i)
		 (<= (arc-head arc) i)
		 ;; This next condition is an optimization that relies
		 ;; on the fact that we are building up the assignments
		 ;; sequentially and so we have already checked all the
		 ;; constraints on the previous variables and so we
		 ;; care only about constraints involving the NEW var.
		 (or *check-all-arcs*	; disables optimization
		     (= (arc-tail arc) i) (= (arc-head arc) i))
		 )
	    (set-new-arcs-for-var! i (cons arc (new-arcs-for-var i)))))
      arcs))
   (csp-variables *csp*))
  (display* "Done"))

;;; Updates the domain of arc-tail so that only values consistent with
;;; SOME value in the domain of arc-head remain.
(define (revise arc)
  ;; Keep track of whether some label in the domain at the tail of the
  ;; arc is eliminated because there is no label in the domain at the
  ;; head of the arc that is consistent with it.
  (let ((tail-var (arc-tail arc))
	(head-var (arc-head arc))
	(constraint (arc-constraint-function arc))
	(deleted-something? #f))
    (cond (*verbose-revise*
	   (display* "Revising " arc "\nDomain of var=" tail-var " starts as")
	   (pretty-print (domain-of tail-var))
	   ))
    (set-var-domain!			
     tail-var				; modifies the domain of tail-var
     ;; collect only values of tail-var consistent with SOME value of head-var.
     (map-non-false
      (domain-of tail-var)
      (lambda (x)
	(if (not
	     (there-exists?		; returns #t if some element is consistent
	      (domain-of head-var)
	      (lambda (y)
		(set! *number-of-arc-tests* (+ 1 *number-of-arc-tests*))
		(constraint tail-var x head-var y))))
	    ;; No consistent value found, so we will flush this value
	    ;; from the domain of tail-var and we will set deleted-something?
	    ;; to indicate that it happened.
	    (begin (set! deleted-something? #t)
		   #f)
	    ;; Some consistent value found, so keep this value.
	    x)
	)))
    (if *verbose-revise*
	(cond (deleted-something?
	       (display* "Domain of var=" tail-var "ends as")
	       (pretty-print (domain-of tail-var)))
	      (else
	       (display* "Domain of var=" tail-var "did not change.")
	   )))
    deleted-something?))

;;; This applies REVISE repeatedly.  A change in the domain of some
;;; variable will cause reexamination of constraints involving that
;;; variable.  This will return #f if an empty domain is found, #t
;;; otherwise. 
(define (propagate-constraints q)
  (define (propagate)
    (if (null? q)
	#t
	(let* ((arc (first q))		; first one on q
	       (tail-var (arc-tail arc))
	       (head-var (arc-head arc))
	       (constraint (arc-constraint-function arc))
	       (fail? #f))
	  (set! q (rest q))		; eliminate arc from q
	  (if (revise arc)		; the domain of tail-var changed
	      ;; Add to q all the arcs of the form (i tail-var), that is all
	      ;; arcs whose second variable is tail-var.  Since the domain of
	      ;; tail-var has changed, it is possible that calling
	      ;; revise(i,tail-var) will change the domain of i.  This is
	      ;; where further constraint propagation happens.
	      (if (null? (domain-of tail-var))
		  (set! fail? #t)
		  (for-each (lambda (a)
			      (if (and (not (= (arc-tail a) head-var))
				       (not (member a q)))
				  (set! q (cons a q))))
			    ;; arcs whose arc-head matches tail-var of
			    ;; the current arc, ie whose form is (i tail-var)
			    (arcs-into tail-var))))
	  (if fail?
	      #f
	      (propagate)))))
  (if *verbose* (print-domain-sizes "Domain sizes before propagation"))
  (let ((answer (propagate)))
    (if *verbose* (print-domain-sizes "Domain sizes after propagation"))
    answer)
  )

(define (length=1? list)
  ;; This checks for a length=1 list without finding length.
  (and (list? list) 
       (not (eq? list '()))
       (eq? (cdr list) '())))

;;; The simple version of propagate-constraints that does no
;;; propagation at all except to those vars directly connected to
;;; input var by a constraint.  Returns #f if a contradiction found.
(define (forward-check var)
  (define (forward-check-loop arcs)	; returns #t or #f
    (if (null? arcs)			; did all arcs with no failure
	#t				; indicate success
	(begin 
	  (revise (first arcs))		; update the domain of tail-var
					; of first arc.
	  (if (null? (domain-of (arc-tail (first arcs))))
	      #f			; failure if domain became empty
	      (forward-check-loop (rest arcs))))))

  (if *verbose* (pretty-print (cons (list 'before var) (get-list-of-domains))))
  (let ((answer
	 ;; update domains of variables constrained by var, i.e.
	 ;; the tail-var of those arcs whose head-var = var.
	 ;; returns #f if a contradiction is found.
	 (forward-check-loop (arcs-into var)) ))
    (if *verbose* (pretty-print (cons (list 'after var) (get-list-of-domains))))
    answer)
  )

;;; This counts how many domain revisions (deletions) a call to
;;; forward-check (after var is set to value) would do , assuming it
;;; didn't stop at conflicts.  This does not change any of the
;;; domains, otherwise the operation is quite analogous to
;;; forward-check.
(define (count-forward-check-deletions var value)
  ;; Accumulate the count for each arc in the input list.
  (define (forward-check-count-loop arcs)
    (if (null? arcs)
	0
	(+ (count-revisions (first arcs) value)
	   (forward-check-count-loop (rest arcs)))))
  ;; look at domains of variables constrained by var, i.e.
  ;; the tail-var of those arcs whose head-var = var.
  (forward-check-count-loop (arcs-into var))
  )

;;; A simple variation of REVISE that only counts revisions, this
;;; assumes it is being used in a forward-check context where the arcs
;;; all share a head-var and the value of that shared variable is a
;;; singleton, namely VALUE.  This function does not change any of
;;; the domains.
(define (count-revisions arc value)
  (let ((tail-var (arc-tail arc))
	(head-var (arc-head arc))
	(constraint (arc-constraint-function arc))
	(deleted-count 0))		; count of deleted domain vals
    (for-each
     (lambda (x)
       (if (constraint tail-var x head-var value) ;check constraint
	   #f				; consistent, do nothing
	   (set! deleted-count (+ 1 deleted-count));; a revision, count it.
	   )
       ;; update the arc test count.
       (set! *number-of-arc-tests* (+ 1 *number-of-arc-tests*)))
     ;; we do the operation above for each value in domain of tail-var.
     (domain-of tail-var))
    ;; return the count of deletions.
    deleted-count))

;;;;;;;;;;;;;;;;;
;;; BACKTRACK
;;;;;;;;;;;;;;;;;

;;; Recursive implementation of backtrack.
(define (backtrack)
  ;; this function does all the work
  (define (backtrack-aux var var-values reversed-partial-answer)
    ;; called when a tentative assignment fails.
    (define (fail)
      (if *verbose* 
	  (display* "Failing for var=" var
		    " value =" (first var-values)))
      ;; try again with a new value for var
      (backtrack-aux var (rest var-values)
		     reversed-partial-answer))

    (if (null? var-values)
	#f				; no values left.
	(let* ((extended-reversed-partial-answer
		(cons (first var-values) reversed-partial-answer)))
	  (display var) (display " ")	; announce the variable
	  (if (consistent? var extended-reversed-partial-answer)
	      ;; consistent partial assignment found
	      (if (= var (- (csp-number-of-variables *csp*) 1))
		  ;; we found a consistent choice for all vars
		  (reverse extended-reversed-partial-answer)
		  ;; try extending answer to next variable, else fail.
		  (or (backtrack-aux (+ var 1) 
				     (domain-of (+ var 1))
				     extended-reversed-partial-answer)
		      (fail)))
	      ;; not consistent, so fail.
	      (fail)))))

  ;; Start backtracking with the first (index = 0) variable.
  (backtrack-aux 0 (domain-of 0) '()))

;; Backtracking with limited constraint propagation (forward-check)
;; after every tentative assignment to a variable.
(define (backtrack-fc)
  ;; this function does all the work
  (define (backtrack-fc-aux var var-values)
    ;; called when a tentative assignment fails.
    (define (fail saved-domains)
      (if *verbose* 
	  (display* "Failing for var=" var
		    " value=" (first var-values)))
      ;; could not extend to next variable, undo side effects
      (restore-domains saved-domains)
      ;; try again with a new value for var
      (backtrack-fc-aux var (rest var-values)))

    (if (null? var-values)		; no values left.
        #f
	(let ((domains (get-list-of-domains)))	; Note
	  (display var) (display " ")	; announce the variable
	  ;; Here we set the domain to have just the one specific value.
	  (set-var-domain! var (list (first var-values)))
	  ;; Propagate the effects of this choice.  Recall fc returns #f
	  ;; if a contradiction is found.
	  (if (forward-check var)
	      ;; End test, assuming numerically ordered variables.
	      (if (= var (- (csp-number-of-variables *csp*) 1))
		  ;; we found a consistent choice for all vars
		  (map car (get-list-of-domains))
		  ;; try extending to another variable, else fail.
		  (or (backtrack-fc-aux (+ var 1) 
					(domain-of (+ var 1)))
		      (fail domains)))
	      ;; not consistent, so fail.
	      (fail domains)))))

  ;; Start backtracking with the first (index = 0) variable.
  ;; This function assumes that the variables are considered in numerical order.
  (backtrack-fc-aux 0 (domain-of 0)))

;; Backtracking with full constraint propagation (propagate-constraints)
;; after every tentative assignment to a variable.
(define (backtrack-pc)
  ;; this function does all the work
  (define (backtrack-pc-aux var var-values)
    ;; called when a tentative assignment fails.
    (define (fail saved-domains)
      (if *verbose* 
	  (display* "Failing for var=" var
		    " value=" (first var-values)))
      ;; could not extend to next variable, undo side effects
      (restore-domains saved-domains)
      ;; try again with a new value for var
      (backtrack-pc-aux var (rest var-values)))

    (if (null? var-values)		; no values left.
        #f
	(let ((domains (get-list-of-domains)))	; Note
	  (display var) (display " ")	; announce the variable
	  ;; Here we set the domain to have just the one specific value.
	  (set-var-domain! var (list (first var-values)))
	  ;; Propagate the effects of this choice.  
	  (propagate-constraints (csp-arcs *csp*))
	  (if (not (contradiction?))
	      (if (= var (- (csp-number-of-variables *csp*) 1))
		  ;; we found a consistent choice for all vars
		  (map car (get-list-of-domains))
		  ;; try extending to another variable, else fail.
		  (or (backtrack-pc-aux (+ var 1) 
					(domain-of (+ var 1)))
		      (fail domains)))
	      ;; not consistent, so fail.
	      (fail domains)))))

  ;; Start backtracking with the first (index = 0) variable.
  (backtrack-pc-aux 0 (domain-of 0)))

;;; Checks whether the reversed-values list represents a valid partial
;;; assignment.  The relevant vars are in the range [0, last-var-index].
(define (consistent? last-var-index reversed-values)
  (for-all? 
   ;; (csp-arcs-added-for-new-var *csp*) holds the constraints involving the
   ;; last-var-index and some other var with a smaller number and
   ;; therefore with a value present in reversed-values.
   (new-arcs-for-var last-var-index)
   (lambda (arc)
     ;; Remember the values are reversed and so we have to index
     ;; backwards.  
     (let ((value-1 (list-ref reversed-values 
			      (- last-var-index (arc-tail arc))))
	   (value-2 (list-ref reversed-values 
			      (- last-var-index (arc-head arc)))))
       (set! *number-of-arc-tests* (+ 1 *number-of-arc-tests*))
       ;; Check the arc constraint on the relevant values.
       ((arc-constraint-function arc)
	(arc-tail arc) value-1 (arc-head arc) value-2)))))

(define (contradiction?)
  (there-exists?
      (csp-variables *csp*)
    (lambda (i) (null? (domain-of i)))))

;;;;;;;;;;;;;;;;;
;;; RUNNING TEST CASES
;;;;;;;;;;;;;;;;;

(define (do-test function)
  (set! *number-of-arc-tests* 0.)
  (let ((answer (function)))
    (print-answer answer)
    (display* "Check-assignment says " (check-assignment answer))
    (print-tests)
    answer))

(define (print-tests)
  (display* "Used " *number-of-arc-tests* " arc tests"))

(define (print-answer answer)
  (if answer
      (for-each
       (lambda (i)
	 (display* "Var: " (name-of i) " Value: " (list-ref answer i)))
       (csp-variables *csp*))
      (display* "No answer found."))
  (newline))

(define (print-domain-sizes . message)
  (if message (display* (car message)))
  (for-each
   (lambda (d) 
     (display* d ": " (length (domain-of d))))
   (csp-variables *csp*))
  (newline))

(define (print-domains . message)
  (if message (display* (car message)))
  (for-each
   (lambda (d) (display* d ":Dom(" (name-of d) ")=" (domain-of d)))
   (csp-variables *csp*))
  (newline))

;;; Checks an assignment (as a list of values for the variables)
;;; against all the constraints in the arcs.
(define (check-assignment answer)
  (let ((check #t))
    (if answer
	(for-all? 
	    (csp-arcs *csp*)
	  (lambda (arc)
	    (let ((value-1 (list-ref answer (arc-tail arc)))
		  (value-2 (list-ref answer (arc-head arc))))
	      (if
	       ((arc-constraint-function arc)
		(arc-tail arc) value-1 (arc-head arc) value-2)
	       #t
	       (begin
		 (display* "Failed constraint " arc value-1 value-2)
		 (set! check #f)
		 #t))))
	  ))
    check))