;;;; -*- 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* (make-vector 0))	; vector of arcs indexed by head-var
(define *arcs-into* (make-vector 0))	; vector of arcs indexed by tail-var
(define *arcs-added-for-new-var* (make-vector 0)) ; vector of arcs, used in checking consistency
(define *variable-domains* (make-vector 0)) ; vector of lists of var values
(define *variable-names* (make-vector 0)) ; 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)
(define *verbose-revise* #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 (get-list-of-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).
  ;;
  (display* "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*)
  (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?))


(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)			; 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) '()))

;;; 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)))))

(define (contradiction?)
  (there-exists?
      *variables*
    (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)))
       *variables*)
      (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))))
   *variables*)
  (newline))

(define (print-domains . message)
  (if message (display* (car message)))
  (for-each
   (lambda (d) (display* d ":Dom(" (name-of d) ")=" (domain-of d)))
   *variables*)
  (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? 
	    *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
		 (display* "Failed constraint " arc value-1 value-2)
		 (set! check #f)
		 #t))))
	  ))
    check))