;; code for red-black dictionary lab

;; The timing procedure, for timing the application of f to the arguments
;;  a1 and a2:

(define time-it
  (lambda (f a1 a2)
    (define runtime
      (lambda ()
        (/ (- (current-process-milliseconds)
              (current-gc-milliseconds))
           1000.)))
    (define loop
      (lambda (count start-time)
        (let ((value (f a1 a2)))
          (let ((cur-time (runtime)))
            (cond ((> cur-time (+ 1 start-time))
                   (newline)
                   (display ";Time: ")
                   (display (/ (- cur-time start-time) count))
                   (newline)
                   value)
                  (else
                   (loop (+ count 1) start-time)))))))
    (loop 1 (runtime))))

;; Some procedures from chapter 7 that will come in handy:

(define filter
  (lambda (ok? lst)
    (cond ((null? lst)
           '())
          ((ok? (car lst))
           (cons (car lst) (filter ok? (cdr lst))))
          (else
           (filter ok? (cdr lst))))))

(define first-elements-of
  (lambda (n list)
    (if (= n 0)
        '()
        (cons (car list) 
              (first-elements-of (- n 1)
                                 (cdr list))))))

;; Unmodified stuff from chapter 13:

(define make-empty-ranked-btree
  (lambda ()
    (let ((tree (make-vector 6)))
      (vector-set! tree 0 #t) ; empty-tree? = true
      (vector-set! tree 2 #f) ; has no parent
      (vector-set! tree 5 0)  ; rank = 0
      tree)))

(define empty-tree?
  (lambda (tree)
    (vector-ref tree 0)))

(define set-empty! ;makes tree empty
  (lambda (tree)
    (vector-set! tree 0 #t)))

(define value
  (lambda (tree)
    (vector-ref tree 1)))

(define set-value!
  (lambda (tree item)
    (vector-set! tree 0 #f) ;not empty
    (vector-set! tree 1 item)))

(define parent
  (lambda (tree)
    (vector-ref tree 2)))

(define root?
  (lambda (tree)
    (not (vector-ref tree 2))))

(define left-subtree
  (lambda (tree)
    (vector-ref tree 3)))

(define set-left-subtree!
  (lambda (tree new-subtree)
    (vector-set! new-subtree 2 tree) ;parent
    (vector-set! tree 3 new-subtree)))

(define right-subtree
  (lambda (tree)
    (vector-ref tree 4)))

(define set-right-subtree!
  (lambda (tree new-subtree)
    (vector-set! new-subtree 2 tree) ;parent
    (vector-set! tree 4 new-subtree)))

(define rank
  (lambda (tree)
    (vector-ref tree 5)))

(define set-rank!
  (lambda (tree rank)
    (vector-set! tree 5 rank)))

(define which-subtree
  (lambda (tree)
    ;; Returns the symbol left if tree is left-subtree of its
    ;; parent and the symbol right if it is the right-subtree
    (cond ((root? tree)
           (error "WHICH-SUBTREE called at root of tree."))
          ((eq? tree (left-subtree (parent tree)))
           'left)
          (else 'right))))

(define sibling
  (lambda (tree)
    (cond ((root? tree)
           (error "SIBLING called at root of tree."))
          ((equal? (which-subtree tree) 'left)
           (right-subtree (parent tree)))
          (else
           (left-subtree (parent tree))))))

(define make-binary-search-tree make-empty-ranked-btree)

(define make-red-black-tree make-binary-search-tree)

(define promote!
  (lambda (node)
    (set-rank! node (+ (rank node) 1))))

(define exchange-values!
  (lambda (node-1 node-2)
    (let ((value-1 (value node-1)))
      (set-value! node-1 (value node-2))
      (set-value! node-2 value-1))))

(define exchange-left-with-right!
  (lambda (tree-1 tree-2)
    (let ((left (left-subtree tree-1))
          (right (right-subtree tree-2)))
      (set-left-subtree! tree-1 right)
      (set-right-subtree! tree-2 left))))

(define rotate-left!
  (lambda (bs-tree)
    (exchange-left-with-right! bs-tree
                               (right-subtree bs-tree))
    (exchange-left-with-right! (right-subtree bs-tree)
                               (right-subtree bs-tree))
    (exchange-left-with-right! bs-tree
                               bs-tree)
    (exchange-values! bs-tree (left-subtree bs-tree))
    'done))

(define rotate-right!
  (lambda (bs-tree)
    (exchange-left-with-right! (left-subtree bs-tree)
                               bs-tree)
    (exchange-left-with-right! (left-subtree bs-tree)
                               (left-subtree bs-tree))
    (exchange-left-with-right! bs-tree
                               bs-tree)
    (exchange-values! bs-tree (right-subtree bs-tree))
    'done))

(define string-comparator
  (lambda (string-1 string-2)
    (cond ((string<? string-1 string-2) '<)
          ((string=? string-1 string-2) '=)
          (else '>))))

(define make-dictionary
  (lambda (key-comparator key-extractor)
    (vector key-comparator
            key-extractor
            (make-red-black-tree))))

(define key-comparator
  (lambda (dictionary)
    (vector-ref dictionary 0)))

(define key-extractor
  (lambda (dictionary)
    (vector-ref dictionary 1)))

(define red-black-tree
  (lambda (dictionary)
    (vector-ref dictionary 2)))

;; The below is an exercise solution:

(define display-ranked-btree
  (lambda (tree)
    (define display-down-from
      (lambda (node depth)
        (display-times "  " depth)
        (display
         (if (empty-tree? node)
             "empty"
             (value node)))
        (display " (rank ")
        (display (rank node))
        (display ")")
        (newline)
        (if (not (empty-tree? node))
            (begin
              (display-down-from (left-subtree node) (+ depth 1))
              (display-down-from (right-subtree node) (+ depth 1))))))
    (display-down-from tree 0)))

(define display-times  ;from chapter 10
  (lambda (output count)
    (if (= count 0)
        'done
        (begin (display output)
               (display-times output (- count 1))))))

;; Now comes some of the modifications called for in 13.6:

(define insertion-point
  (lambda (item bs-tree key-comparator key-extractor)
    ;; This procedure finds the point at which item should be
    ;; inserted in bs-tree. In other words, it finds the empty
    ;; leaf node where it should be inserted so that the
    ;; binary search condition still holds after it is inserted.
    ;; If item is already in bs-tree, then the insertion
    ;; point will be found by searching to the right so that
    ;; the new copy will occur "later" in bs-tree.
    (cond ((empty-tree? bs-tree) bs-tree)
          (else
           (let ((c (key-comparator (key-extractor item)
                                    (key-extractor (value bs-tree)))))
             (cond ((equal? c '<)
                    (insertion-point item (left-subtree bs-tree)
                                     key-comparator key-extractor))
                   (else
                    (insertion-point item (right-subtree bs-tree)
                                     key-comparator key-extractor))))))))

(define binary-search-insert!
  (lambda (item bs-tree key-comparator key-extractor)
    ;; This procedure will insert item into bs-tree at a leaf
    ;; (using the procedure insertion-point), maintaining
    ;; the binary search condition on bs-tree. The return value
    ;; is the subtree that has item at its root.
    ;; If item occurs in bs-tree, another copy of item
    ;; is inserted into bs-tree
    (let ((insertion-tree (insertion-point item bs-tree
                                           key-comparator key-extractor)))
      (set-value! insertion-tree item)
      (set-left-subtree! insertion-tree
                         (make-binary-search-tree))
      (set-right-subtree! insertion-tree
                          (make-binary-search-tree))
      insertion-tree)))

(define red-black-insert!
  (lambda (item red-black-tree key-comparator key-extractor)
    (define rebalance!
      (lambda (node)
        (cond ((root? node)
               'done)
              ((root? (parent node))
               'done)
              ((< (rank node) (rank (parent (parent node))))
               'done)
              ((= (rank node) (rank (sibling (parent node))))
               (promote! (parent (parent node)))
               (rebalance! (parent (parent node))))
              (else
               (let ((path-from-grandparent
                      (list (which-subtree (parent node))
                            (which-subtree node))))
                 (cond ((equal? path-from-grandparent '(left left))
                        (rotate-right! (parent (parent node))))
                       ((equal? path-from-grandparent '(left right))
                        (rotate-left! (parent node))
                        (rotate-right! (parent (parent node))))
                       ((equal? path-from-grandparent '(right left))
                        (rotate-right! (parent node))
                        (rotate-left! (parent (parent node))))
                       (else ; '(right right)
                        (rotate-left! (parent (parent node))))))))))
    (let ((insertion-node (binary-search-insert! item red-black-tree
                                                 key-comparator
                                                 key-extractor)))
      (set-rank! insertion-node 1)
      (rebalance! insertion-node))
    'done))

;; We want a variant of string-comparator that deals with *.  This can
;; be done by truncating the * off the end of the string containing it
;; and truncating the other string to the same length, and then comparing
;; with string-comparator.  We can do the truncating with substring.  The
;; string-pos procedure, defined below, is used to find the position of the
;; *, or #f if there is no *.

(define wildcarded-string-comparator
  (lambda (s1 s2)
    (let ((p1 (string-pos s1 #\*)))
      (if p1
          (wildcarded-string-comparator
           (substring s1 0 p1)
           (substring s2 0 (min p1 (string-length s2))))
          (let ((p2 (string-pos s2 #\*)))
            (if p2
                (wildcarded-string-comparator
                 (substring s1 0 (min p2 (string-length s1)))
                 (substring s2 0 p2))
                (string-comparator s1 s2)))))))

(define string-pos
  (lambda (s c) ;find index in s of first c, of #f if none
    (define loop
      (lambda (i)
        (if (= i (string-length s))
            #f
            (if (char=? (string-ref s i) c)
                i
                (loop (+ i 1))))))
    (loop 0)))

;; Here we load in some sample data, of the form
;;   (define names
;;      '(("first-name" "last-name" "login-name")
;;        ("another-first" "another-last" "another-login")
;;        ...))
;; Since the list is thousands of elements long, I didn't
;; want to include it here.
;;
;; The data is for use only in a lab exercise, and should not
;; be redistributed or otherwise used.

(load "~mc28/labs/red-black-dictionaries/names.scm")

;; Now we can make a dictionary for holding this data, indexed by last-name
;; (which is in the cadr, i.e. second list element -- we're not bothering
;;  with an ADT).

(define names-by-last-name
  (make-dictionary wildcarded-string-comparator cadr))

;; It is convenient to have a mutator that totally cleans out a dictionary,
;; so that in doing timing trials you can clean out the dictionary and then
;; refill it (using for-each and dictionary-insert!) with the number of entries
;; you want to use, which you can get from the names list using the
;; first-elements-of procedure.

(define reset-dictionary!
  (lambda (dictionary)
    (vector-set! dictionary 2 (make-red-black-tree))))

;; For comparison with dictionary-retrieve, here is list-retrieve, which
;; has the property that (list-retrieve "WildcardedLastName" names) should
;; retrieve the same list of matches (though possibly in a different order)
;; as (dictionary-retrieve "WildcardedLastName" names-by-last-name), assuming
;; all the entries in names have been inserted into names-by-last-name.
;; WildcardedLastName can be a last name, like Hailperin, or a last name
;; prefix with * at the end, like Hai*

(define list-retrieve
  (lambda (name lst)
    (filter (lambda (entry)
              (equal? (wildcarded-string-comparator (cadr entry) name)
                      '=))
            lst)))

;; At this point, the following tasks remain to be done for the lab:
;;
;; (1) You should write the dictionary-insert! procedure and test it
;;     by inserting some entries (such as the first few from names)
;;     into names-by-last-name and then displaying the dictionary's
;;     red-black tree by doing
;;       (display-ranked-btree (red-black-tree names-by-last-name))
;;     Note that the dictionary-insert! procedure should not be
;;     large or complex: it just calls red-black-insert! with the
;;     appropriate arguments.
;;
;; (2) Since you'll want to do tests with varying numbers of names
;;     in the red-black tree, you should write a procedure that takes
;;     an argument specifying the number of names, n, to be used and
;;     does a reset-dictionary! on names-by-last-name and then inserts
;;     the first n elements of names into names-by-last-name using
;;     the dictionary-insert! procedure you wrote, in conjunction with
;;     for-each (as described in section 13.6) and the first-elements-of
;;     procedure.
;;
;; (3) Write binary-search-retrieve as described in 13.6, and then
;;     define red-black-retrieve as identical to it, i.e.
;;      (define red-black-retrieve binary-search-retrieve)
;;     and then define dictionary-retrieve to call red-black-retrieve
;;     on the dictionary's red-black tree using the dictionary's
;;     key comparator and extractor.  Test this using names-by-last-name.
;;     You should be able to, for example, find all the users with your
;;     last name, or find all users with last names of the form "Hans*",
;;     i.e., starting with Hans.
;;
;; (4) Now do timing tests comparing dictionary-retrieve with list-retrieve
;;     as the number of entries being searched is increased.  You will
;;     get the most distinct results if the search is one that matches
;;     very few names.  However, you may also want to experimentally see
;;     what effect the number of matches has on the time for each of the two
;;     retrieval methods.