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

;;; CSP solution.

;;;;;;;;;;;;;;;;;
;;; DATA STRUCTURES AND ABSTRACTIONS
;;;;;;;;;;;;;;;;;

;;; Global variables that hold the current constraint graph.

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

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

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

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

;;; Arcs are of the form (tail-var-index head-var-index constraint-function index)
(define make-arc list)
(define arc-tail first)
(define arc-head second)
(define arc-constraint-function third)
(define arc-index fourth)

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

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

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

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

;;; Operations on all the domains.

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

;;(define (restore-domains saved)
;;  (set! *variable-domains* (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 *number-of-variables*)
	#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).
  ;;
  (format #t "~%Initializing...")
  (set! *arcs* arcs)
  (set! *number-of-variables* (length names-and-domains))
  (set! *variables* (make-index-list *number-of-variables*))
  (set! *variable-names* (list->vector (map first names-and-domains)))
  (set! *variable-domains* (list->vector (map rest names-and-domains)))
  (set! *arcs-from* (make-vector *number-of-variables*))
  (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! *arcs-into* (make-vector *number-of-variables*))
  (for-each
   (lambda (arc)
     (set-arcs-into! (arc-head arc) 
		     (cons arc (arcs-into (arc-head arc)))))
   arcs)
  ;; The *arcs-added-for-new-var* 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! *arcs-added-for-new-var* (make-vector *number-of-variables*))
  (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))
   *variables*)
  (format #t "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))
    (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 the first true one or #f
	      (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)
	)))
    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)
    (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* (pp (cons (list 'before var) (list-domains))))
  (let ((answer
	 ;; update domains of variables constrained by var, i.e.
	 ;; the tail-var of those arcs whose head-var = var.
	 (forward-check-loop (arcs-into var)) ))
    (if *verbose* (pp (cons (list 'after var) (list-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* 
	  (format #t "Failing for var = ~a value = ~a~%"
		  var (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)))
	  (format #t "~s " var)		; announce the variable
	  (if (consistent? var extended-reversed-partial-answer)
	      ;; consistent partial assignment found
	      (if (= var (- *number-of-variables* 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* 
	  (format #t "Failing for var = ~a value = ~a~%"
		  var (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 (list-domains)))	; Note
	  (format #t "~s " var)		; 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)
	      (if (= var (- *number-of-variables* 1))
		  ;; we found a consistent choice for all vars
		  (map car (list-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.
  (backtrack-fc-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? 
   ;; *arcs-added-for-new-var* 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)))))

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

(define (do-test function)
  (set! *number-of-arc-tests* 0.)
  (let ((answer (function)))
    (print-answer answer)
    (format #t "~%Check-assignment says ~s" (check-assignment answer))
    (print-tests)
    answer))

(define (print-tests)
  (format #t "~%Used ~s arc tests~%" *number-of-arc-tests*))

(define (print-answer answer)
  (if answer
      (for-each
       (lambda (i)
	 (format #t "Var: ~s Value: ~s~%" (name-of i) (list-ref answer i)))
       *variables*)
      (format #t "No answer found."))
  (newline))

(define (print-domain-sizes message)
  (format #t "~a: " message)
  (for-each
   (lambda (i) (format #t "~s " (length (domain-of i))))
   *variables*)
  (newline))

;;; Checks an assignment (as a list of values for the variables)
;;; against all the constraints in the arcs.
(define (check-assignment answer)
  (if answer
      (for-all? 
       *arcs*
       (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
	      (format #t "Failed constraint ~s ~s ~s~%"
		      arc value-1 value-2)
	      #f))))
       )))