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

;;;; SOME GLOBAL VARIABLES

(define *nn-verbose* #f)
(define *kd-tree* #f)

(define *node-type-id* 0)
(define *node-dimension-id* 1)
(define *node-gap-min-id* 2)
(define *node-gap-max-id* 3)
(define *node-left-branch-id* 4)
(define *node-right-branch-id* 5)

;;; Controls drawing
(define *draw* #f)		; turn on display?

(define *neighbors* 1)			; The K in k-nearest-neighbors

;;; High-Level interface for classification

;; The actual names of the classes
(define *nn-class0* #f)			; set by nn-train
(define *nn-class1* #f)

(define (nn-train data)
  (let ((classes (map car (class-counts data))))
    (cond ((= (length classes) 2)
	   (cond ((or (equal? classes '(0 1))
		      (equal? classes '(0.0 1.0))
		      (equal? classes '(1 0))
		      (equal? classes '(1.0 0.0)))
		  (set! *nn-class0* 0)
		  (set! *nn-class1* 1))
		 (else
		  (set! *nn-class0* (first classes))
		  (set! *nn-class1* (second classes)))))
	  (else 
	   (error "Only two classes allowed, but we have" classes))))
  (make-kd-tree data)
  'done)

(define (nn-classify sample) (nn-identify sample))

(define (make-nn-prediction unknown)
  (let ((candidates (find-nn-answers unknown *kd-tree* *neighbors* 0)))
    (if (= 0 (nn-answer-distance^2 (first candidates)))
	;; perfect match
	winner
	;;Otherwise, tally up votes:
	(first (tally-votes candidates *weighting-function*)))))

;;;; DATA STRUCTURES

(define-structure (node)
  dimension				; The dimension that the node divides on.
  gap-min				; The maximum of the "small" samples.
  gap-max				; The minimum of the "large" samples.
  left-branch				; The left sample(s) or node.
  right-branch				; The right sample(s) or node.
  )

(define-structure (nn-answer)
  diagnosis distance^2 record)

;;;; CONSTRUCT KD-TREE

(define *maximum-depth* 0)

;;  Purpose:	Supply arguments to make-kd-tree-aux and assign value to *kd-tree*.
;;  Returns:	A kd-tree node (a structure).
(define (make-kd-tree data)
  (set! *maximum-depth* 0)
  (set! *kd-tree* 
	(make-kd-tree-aux
	 data
	 (length (data-point-features (first data)))	; number dimensions
	 0
	 0
	 ))
  (display* "The maximum depth of the kd-tree is " *maximum-depth*)
  (if *draw* 
      (draw-classifier
       nn-get-point
       (map (lambda (p) 
	      (make-data-point (list (if (eq? (data-point-class p) *nn-class0*) 0 1)) 
			       (data-point-features p)))
	    data)))
  #t)

;;  Purpose:	Construct a KD kd-tree from samples.
;;  Returns:	A kd-tree node (a structure).
(define (make-kd-tree-aux samples dim-count level skipped)
  (set! *maximum-depth* (max *maximum-depth* level))
  (let ((dimension-to-check (remainder level dim-count)))
    (cond
     ((null? (cdr samples))
      (first samples))
     ((>= skipped dim-count) 
      ;; found a bunch of identical sample
      (cond ((= (length (class-counts samples)) 1)
	     (first samples))
	    (else
	     (display* "These samples differ in class and all have same features.
Creating a sample with majority class.")
	     (for-each display* samples)
	     ;; create an artificial point with the majority class
	     (cons (list (most-common-class samples))
		   (data-point-features (first samples)))))
      )
     ((variation-in-dimension? dimension-to-check samples)
      (let ((node (make-node-for-samples dimension-to-check samples)))
	;;Recurse along both branches:
	(set-node-left-branch
	 node
	 (make-kd-tree-aux (node-left-branch node) dim-count (1+ level) 0))
	(set-node-right-branch
	 node
	 (make-kd-tree-aux (node-right-branch node) dim-count (1+ level) 0))
	node))
     (else 
      (if *nn-verbose*
	  (display*
	   "Skipping dimension "
	   dimension-to-check
	   " on level "
	   level
	   "- no variation"))
      (make-kd-tree-aux samples dim-count (1+ level) (+ skipped 1))))))

;;  Purpose:	See if the samples vary in the given dimension.
;;  Returns:	#f or #t
(define (variation-in-dimension? dimension samples)
  (let ((reference (list-ref (data-point-features (first samples)) dimension)))
    (define (loop others)
      (if (null? others)
	  #f
	  (if (equal? reference
		      (list-ref (data-point-features (first others)) dimension))
	      (loop (cdr others))
	      #t)))
    (if (null? (cdr samples))
	#f
	(loop (cdr samples)))))

;;  Purpose:	Sorts samples along dimension supplied and creates a node.
;;  Returns:	A node.
(define (make-node-for-samples dimension samples)
  (let* ((samples
	  ;; Sort the samples:
	  (sort samples
		(lambda (x y)
		  (< (list-ref (data-point-features x) dimension)
		     (list-ref (data-point-features y) dimension)))))
	 (numbers
	  ;; Extract numbers in given dimension:
	  (map (lambda (x) (list-ref (data-point-features x) dimension))
	       samples))
	 (index-limit 
	  ;; Compute the position of the rightmost element of the list:
	  (-1+ (length numbers)))
	 (left-size
	  ;; Compute the position of the center of the list:
	  (quotient (length numbers) 2)))
    ;; Purpose:	Finds nearest pair of samples that differ in the dimension.
    ;;          Sets left-size variable accordingly.
    (define (seek-nearest-split n)
      (let* ((ll (max 0 (-1+ (- left-size n))))
	     (lr (1+ ll))
	     (rr (min index-limit (+ left-size n)))
	     (rl (-1+ rr)))
	(cond ((or (not (<= 0 ll index-limit))
		   (not (<= 0 lr index-limit))
		   (not (<= 0 rl index-limit))
		   (not (<= 0 rr index-limit)))
	       (display* "Limit problem:")
	       (pretty-print (list ll lr rl rr))
	       (pretty-print numbers)
	       (error "Limit problem"))
	      ((not (= (list-ref numbers ll) (list-ref numbers lr)))
	       (set! left-size (1+ ll)))
	      ((not (= (list-ref numbers rl) (list-ref numbers rr)))
	       (set! left-size (1+ rl)))
	      (else 
	       (seek-nearest-split (1+ n))))))
    ;; Reset left-size, if necessary:
    (seek-nearest-split 0)
    ;; Split up samples and make a node:
    (make-node
     dimension
     (list-ref numbers (-1+ left-size))
     (list-ref numbers left-size)
     (first-n left-size samples)
     (last-n left-size samples))))

;;;; DISPLAY KD-TREE

;;  Purpose:	Supply argument to show-kd-tree-aux.
(define (show-kd-tree)
  (show-kd-tree-aux *kd-tree* 0 #f))

;;  Purpose:	Display a kd-tree using indentation to indicate level.
;;  Arguments:	The root node of the kd-tree.
(define (show-kd-tree-aux node level branch)
  (if (node? node)
      (display* 
	      (indent level)
	      (if branch branch "Top")
	      " split on dimension "
	      (node-dimension node)
	      "["
	      (node-gap-min node)
	      (node-gap-max node)
	      "]")
      (display* (indent level) node))
  (if (node? node)
      (begin
	(show-kd-tree-aux
	 (node-left-branch node) (1+ level) "Left branch")
	(show-kd-tree-aux
	 (node-right-branch node) (1+ level) "Right branch"))))

;;;; NUMERICAL AUXILIARIES

;;  Purpose:	Computes squared weighted distance between two vectors.
(define (distance^2 u v)
  (reduce + 0
	  (map (lambda (x y) (delta^2 x y))
	       u v)))

;;  Purpose:	Computes squared weighted distance between two points.
(define (delta^2 x1 x2)
  (let ((delta (- x1 x2)))
    (* delta delta)))

;;;; IDENTIFICATION

;;  Purpose:	Weight neighbors evenly no matter what the distance.
(define (weight-evenly x)
  1)

;;  Purpose:	Weight neighbors inversely proportional to distance (recall x is dist^2).
(define (weight-by-inverse-distance x)
  (/ (sqrt x)))

(define *weighting-function* weight-evenly)

(define (nn-identify unknown)
  (nn-identify-aux unknown *kd-tree* *neighbors* *weighting-function*))

;;  Purpose:	To guess an attribute using nearest neighbor idea.
;;  Arguments:	unknown: a list of dimension values
;;		kd-tree: a KD tree.
;;		count: the number of nearest neighbors to be used.
;;		weighting-function: a function that establishes the way
;;				    a neighbor's influence is diminished
;;				    by distance.
;;  Returns:	The best guess for an unknown's diagnosis.
(define (nn-identify-aux unknown kd-tree count weighting-function)
  (let ((candidates (find-nn-answers (data-point-features unknown) kd-tree count 0)))
    (cond (*nn-verbose*
	   (newline) (display "Unknown:") 
	   (newline) (display unknown)
	   (newline) (display "Candidate(s):")
	   (for-each (lambda (x y)
		       (newline) (display x)
		       (display " ") (display (sqrt y)))
		     (map nn-answer-record candidates)
		     (map nn-answer-distance^2 candidates))
	   (newline)))
    (if (= 0 (nn-answer-distance^2 (first candidates)))
	;;If the answer is at the same place exactly, report it:
	(let ((winner (nn-answer-diagnosis (first candidates))))
	  (display* "The winner is "  winner " (exact match).")
	  winner)
	;;Otherwise, tally up votes:
	(let* ((best-pair (tally-votes candidates weighting-function))
	       (winner (first best-pair))
	       (score (second best-pair)))
	  (display* "The winner is " winner " with " score " votes."
		    "  Correct is " (data-point-class unknown))
	  winner))))

;;  Purpose: Combines evidence when there are multiple nearest neighbors.
;;  Returns: Attribute with the most votes.
(define (tally-votes attribute-distance-pairs weighting-function)
  (let* ((attribute-weight-pairs
	  ;;Make a-list pairs in which first element is an object's attribute 
	  ;;and the second element is the distance-determined influence of
	  ;;that object attribute.
	  (map (lambda (e)
		 (list (nn-answer-diagnosis e)
		       (weighting-function (nn-answer-distance^2 e))))
	       attribute-distance-pairs))
	 (attribute-score-pairs 
	  ;;Make a-list pairs in which first element is an attribute and the
	  ;;second element is the sum of the influences of the objects
	  ;;with that attribute:
	  (map (lambda (attribute)
		 (list attribute
		       (accumulate-weight attribute attribute-weight-pairs)))
	       (remove-duplicates (map first attribute-weight-pairs)))))
    ;;Sort, with the most recommended attribute in front:
    (set! attribute-score-pairs
	  (sort attribute-score-pairs
		(lambda (x y) (> (second x) (second y)))))
    (if *nn-verbose*
	(display* "The scores are: " attribute-score-pairs))
    ;;Pick the winner off the front:
    (first attribute-score-pairs)))

;;  Purpose:	Helper
(define (accumulate-weight attribute pairs)
  (cond ((null? pairs) 0)
	((eq? attribute (first (first pairs)))
	 (+ (second (first pairs)) (accumulate-weight attribute (cdr pairs))))
	(else (accumulate-weight attribute (cdr pairs)))))

;;;; FIND NEAREST NEIGHBORS

;;  Purpose:	Find N nearest neighbors using KD tree.
;;  Returns:	N answers
;;  Remarks:	This procedure is complicated, in part, because
;;		it has to deal with multiple nearest neighbors.
(define (find-nn-answers features kd-tree count level) 

  (define (answers-from-branch branch)
    ;;If the branch is a node, then find the closest neighbor by
    ;;calling find-nn-answers recursively; otherwise the branch is not a
    ;;node, and the branch is the closest neighbor, so return list of that.
    (if (node? branch)
	(find-nn-answers features branch count (1+ level))
	(list (make-nn-answer
	       (data-point-class branch)
	       (distance^2 features
			   (data-point-features branch))
	       branch))))

  (define (combine-answers ans1 ans2)
    ;;Append answers and sort:
    (sort (append ans1 ans2)
	  (lambda (x y) (< (nn-answer-distance^2 x) (nn-answer-distance^2 y)))))

  (let* ((dimension (node-dimension kd-tree))
	 (projection (list-ref features dimension))
	 (left-delta^2 (delta^2 projection
				(node-gap-min kd-tree)))
	 (right-delta^2 (delta^2 projection
				 (node-gap-max kd-tree)))
	 ;;Decide which branch has won and set variables accordingly:
	 (direction (if (< right-delta^2 left-delta^2) 'right 'left))
	 (threshold-delta^2 (if (eq? direction 'right)
				left-delta^2 right-delta^2))
	 (winning-branch
	  (if (eq? direction 'right)
	      (node-right-branch kd-tree) (node-left-branch kd-tree)))
	 (losing-branch 
	  (if (eq? direction 'right)
	      (node-left-branch kd-tree) (node-right-branch kd-tree))))
    (if *nn-verbose*			; report
	(report-direction direction level dimension projection kd-tree))
    ;;At this point, it looks like the winning direction is known.
    ;;This may prove wrong later, of course, because the decision
    ;;is based on comparison in one dimension only, not on actual distance.
    (let* ((winning-answers (answers-from-branch winning-branch))
	   (n-winners (length winning-answers))
	   (nearest-winning-distance^2
	    (nn-answer-distance^2 (first (last-pair winning-answers)))))
      ;;at this point, find-nn-answers needs to check the nearest answer
      ;;by comparing the actual distance between the unknown features and
      ;;the nearest answer with the one-dimensional distance between
      ;;the unknown features and the nearest answer on the wrong side of
      ;;the gap between the left and right groups; also, there
      ;;may not yet be enough answers:
      (first-n				; only keep up to count answers
       count
       (cond ((and (<= nearest-winning-distance^2 threshold-delta^2)
		   (>= n-winners count))
	      ;;If the answer holds up, done:
	      winning-answers)
	     ;;Otherwise, find the best answers on the other side too:
	     (else
	      ;;Indicate why there is more work to do:
	      (if *nn-verbose*
		  (report-continue-reason
		   level nearest-winning-distance^2 threshold-delta^2
		   winning-answers count))
	      ;;Return the combined answers.
	      (combine-answers
	       ;; the answers from the winning branch
	       winning-answers
	       ;;Get best answers on the losing branch of the kd-tree:
	       (answers-from-branch losing-branch)))
	     )))))

(define (report-direction direction level dimension projection kd-tree)
  (cond ((eq? direction 'right)
	 (display*
	  (indent level)
	  "Turn toward large numbers at level " level "(" dimension ")":
	  projection " is closer to " (node-gap-max kd-tree)
	  " than to " (node-gap-min kd-tree)))
	(else
	 (display*
	  (indent level)
	  "Turn toward small numbers at level " level "(" dimension ")":
	  projection " is closer to " (node-gap-max kd-tree)
	  " than to " (node-gap-min kd-tree)))))

(define (report-continue-reason 
	 level nearest-winning-distance^2 threshold-delta^2
	 winning-answers count)
  (if (<= nearest-winning-distance^2 threshold-delta^2)
      (display*
       (indent level)
       "Trying alternate branch because too few answers "
       "[" (length winning-answers) " < " count "]")
      (display*
       (indent level)
       "Trying other branch at level "
       level
       " because worst answer is not good enough " 
       "[" nearest-winning-distance^2 " > " threshold-delta^2 "]")))