; ternary search tries

(define (node v? v s l e h) (vector v? v s l e h))
(define nil (vector #f #f (integer->char 0) (vector) (vector) (vector)))
(define (nil? tst) (eqv? tst nil))
(define (val? tst)  (vector-ref tst 0))
(define (val tst)   (vector-ref tst 1))
(define (split tst) (vector-ref tst 2))
(define (lokid tst) (vector-ref tst 3))
(define (eqkid tst) (vector-ref tst 4))
(define (hikid tst) (vector-ref tst 5))

(define (lookup t k)
  (cond ((nil? t) #f)
        ((null? k) (if (val? t) (val t) #f))
        ((char<? (car k) (split t)) (lookup (lokid t) k))
        ((char<? (split t) (car k)) (lookup (hikid t) k))
        (else (lookup (eqkid t) (cdr k)))))

(define (insert t k v)
  (cond ((nil? t) (insert (node #f #f (if (null? k) (integer->char 0) (car k)) nil nil nil) k v))
        ((null? k) (node #t v (split t) (lokid t) (eqkid t) (hikid t)))
        ((char<? (car k) (split t)) (node (val? t) (val t) (split t) (insert (lokid t) k v) (eqkid t) (hikid t)))
        ((char<? (split t) (car k)) (node (val? t) (val t) (split t) (lokid t) (eqkid t) (insert (hikid t) k v)))
        (else (node (val? t) (val t) (split t) (lokid t) (insert (eqkid t) (cdr k) v) (hikid t)))))

(define (update t k p v)
  (cond ((nil? t) (update (node #f #f (if (null? k) (integer->char 0) (car k)) nil nil nil) k p v))
        ((null? k) (if (val? t) (node #t (p k (val t)) (split t) (lokid t) (eqkid t) (hikid t))
                                (node #t v (split t) (lokid t) (eqkid t) (hikid t))))
        ((char<? (car k) (split t)) (node (val? t) (val t) (split t) (update (lokid t) k p v) (eqkid t) (hikid t)))
        ((char<? (split t) (car k)) (node (val? t) (val t) (split t) (lokid t) (eqkid t) (update (hikid t) k p v)))
        (else (node (val? t) (val t) (split t) (lokid t) (update (eqkid t) (cdr k) p v) (hikid t)))))

(define (delete t k)
  (cond ((nil? t) t)
        ((null? k) (node #f #f (split t) (lokid t) (eqkid t) (hikid t)))
        ((char<? (car k) (split t)) (node (val? t) (val t) (split t) (delete (lokid t) k) (eqkid t) (hikid t)))
        ((char<? (split t) (car k)) (node (val? t) (val t) (split t) (lokid t) (eqkid t) (delete (hikid t) k)))
        (else (node (val? t) (val t) (split t) (lokid t) (delete (eqkid t) (cdr k)) (hikid t)))))

(define (enlist t)
  (let enlist ((t t) (k '()))
    (if (nil? t) '()
      (append (enlist (lokid t) k)
              (if (val? t)
                  (cons (cons (list->string (reverse k)) (val t))
                        (enlist (eqkid t) (cons (split t) k)))
                  (enlist (eqkid t) (cons (split t) k)))
              (enlist (hikid t) k)))))

(define (take n xs)
  (let loop ((n n) (xs xs) (ys '()))
    (if (or (zero? n) (null? xs))
        (reverse ys)
        (loop (- n 1) (cdr xs)
              (cons (car xs) ys)))))

(define (read-word)
  (let loop ((c (read-char)) (cs '()))
    (cond ((eof-object? c)
            (if (null? cs) c
              (list->string (reverse cs))))
          ((char-alphabetic? c)
            (loop (read-char) (cons (char-downcase c) cs)))
          ((pair? cs) (list->string (reverse cs)))
          (else (loop (read-char) cs)))))

(define (word-freq n file-name)
  (define (freq-gt? a b) (> (cdr a) (cdr b)))
  (with-input-from-file file-name
    (lambda ()
      (let loop ((word (read-word)) (freqs nil))
        (if (eof-object? word)
            (take n (sort freq-gt? (enlist freqs)))
            (loop (read-word) (update freqs (string->list word) (lambda (k v) (+ v 1)) 1)))))))

(display (word-freq 25 "bible.txt"))
