;;; Provide Huffman encoding
(in-package "ACL2")

(local (include-book "tools/mv-nth" :dir :system))

;;; Deine huffman tree
(defun nodep (x)
  (declare (xargs :guard t))
  (and (consp x)
       (eq (car x) 'node)))

(defun huffman-treep (x)
  (declare (xargs :guard t))
  (flet ((symbol-list (x)
           (if (nodep x)
               (second x)
               (list (second x)))))
    (if (nodep x)
        (and (true-listp x)
             (equal (len x) 5)
             (true-listp (second x))
             (rationalp (third x))
             (huffman-treep (fourth x))
             (huffman-treep (fifth x))
             (equal (second x)
                    (append (symbol-list (fourth x))
                            (symbol-list (fifth x))))
             (equal (third x)
                    (+ (third (fourth x))
                       (third (fifth x)))))
        (and (true-listp x)
             (equal (len x) 3)
             (equal (first x) 'leaf)
             (rationalp (third x))))))

(defun symbol-list (x)
  (declare (xargs :guard (huffman-treep x)))
  (if (nodep x)
      (second x)
      (list (second x))))

(defun weight (x)
  (declare (xargs :guard (huffman-treep x)))
  (third x))

(defun huffman-leaf (s w)
  (declare (xargs :guard (rationalp w)))
  (list 'leaf s w))

(defun leaf-symbol (x)
  (declare (xargs :guard (and (huffman-treep x)
                              (not (nodep x)))))
  (second x))

(defun huffman-node (l r)
  (declare (xargs :guard (and (huffman-treep l)
                              (huffman-treep r))))
  (list 'node
        (append (symbol-list l)
                (symbol-list r))
        (+ (weight l) (weight r))
        l
        r))

(defun left (x)
  (declare (xargs :guard (and (huffman-treep x)
                              (nodep x))))
  (fourth x))

(defun right (x)
  (declare (xargs :guard (and (huffman-treep x)
                              (nodep x))))
  (fifth x))

(defthm leaf-symbol-weight-elim
  (implies (and (not (nodep x))
                (huffman-treep x))
           (equal (huffman-leaf (leaf-symbol x)
                                (weight x))
                  x))
  :rule-classes :elim)

(defthm len=5
  (implies (true-listp x)
           (equal (equal (len x) 5)
                  (and (consp x)
                       (consp (cdr x))
                       (consp (cddr x))
                       (consp (cdddr x))
                       (consp (cddddr x))
                       (null (cdr (cddddr x)))))))

(defthm left-right-elim
  (implies (and (nodep x)
                (huffman-treep x))
           (equal (huffman-node (left x) (right x))
                  x))
  :hints (("Goal" :use len=5
                  :in-theory (disable len=5)))
  :rule-classes :elim)

(defthm huffman-treep-huffman-node
  (implies (and (huffman-treep l)
                (huffman-treep r))
           (huffman-treep (huffman-node l r))))

(defthm nodep-huffman-node
  (implies (and (huffman-treep l)
                (huffman-treep r))
           (nodep (huffman-node l r))))

(defthm symbol-list-left-right
  (implies (and (nodep x)
                (huffman-treep x))
           (equal (symbol-list x)
                  (append (symbol-list (left x))
                          (symbol-list (right x))))))

(defthm weight-left-right
  (implies (and (nodep x)
                (huffman-treep x))
           (equal (weight x)
                  (+ (weight (left x))
                     (weight (right x))))))

(defthm huffman-treep-left
  (implies (and (nodep x)
                (huffman-treep x))
           (huffman-treep (left x))))

(defthm huffman-treep-right
  (implies (and (nodep x)
                (huffman-treep x))
           (huffman-treep (right x))))

(defthm true-listp-symbol-list
  (implies (huffman-treep x)
           (true-listp (symbol-list x))))

(defthm symbol-list-leaf-symbol
  (implies (and (huffman-treep x)
                (not (nodep x)))
           (equal (symbol-list x)
                  (list (leaf-symbol x)))))

(defthm rationalp-weight
  (implies (huffman-treep x)
           (rationalp (weight x))))

(defthm rationalp-weight/leaf
  (implies (huffman-treep (huffman-leaf s w))
           (rationalp w)))

(defthm huffman-treep-huffman-leaf
  (implies (rationalp w)
           (huffman-treep (huffman-leaf s w))))

(defun node-count (x)
  (declare (xargs :guard (huffman-treep x)))
  (if (nodep x)
      (+ 1
         (node-count (left x))
         (node-count (right x)))
      0))

(defthm node-count-left
  (implies (nodep x)
           (< (node-count (left x))
              (node-count x)))
  :rule-classes :linear)

(defthm node-count-right
  (implies (nodep x)
           (< (node-count (right x))
              (node-count x)))
  :rule-classes :linear)

(in-theory (disable nodep huffman-treep
                    leaf-symbol
                    huffman-leaf huffman-node
                    symbol-list weight left right
                    node-count))

;;; Generate huffman trees
(defun symbol-weight-pair-listp (x)
  (declare (xargs :guard t))
  (if (atom x)
      (null x)
      (and (consp (car x))
           (consp (cdar x))
           (rationalp (cadar x))
           (null (cddar x))
           (symbol-weight-pair-listp (cdr x)))))

(defun huffman-tree-listp (x)
  (declare (xargs :guard t))
  (if (atom x)
      (null x)
      (and (huffman-treep (car x))
           (huffman-tree-listp (cdr x)))))

(defun orderdp (x)
  (declare (xargs :guard (huffman-tree-listp x)))
  (cond ((atom x) t)
        ((<= (weight (car x)) (weight (car x)))
         (orderdp (cdr x)))
        (t nil)))

(defun insert (e x)
  (declare (xargs :guard (and (huffman-treep e)
                              (huffman-tree-listp x))))
  (cond ((endp x) (list e))
        ((<= (weight e) (weight (car x)))
         (cons e x))
        (t
         (cons (car x)
               (insert e (cdr x))))))

(defun pairs-to-leaves (x)
  (declare (xargs :guard (symbol-weight-pair-listp x)))
  (if (endp x)
      nil
      (insert (huffman-leaf (caar x) (cadar x))
              (pairs-to-leaves (cdr x)))))

(defun generate (x a)
  (declare (xargs :measure (len x)
                  :guard (and (huffman-tree-listp x)
                              (orderdp x)
                              (huffman-treep a))))
  (cond ((endp x) a)
        ((<= (weight a) (weight (car x)))
         (generate (cdr x) (huffman-node a (car x))))
        ((endp (cdr x))
         (huffman-node (car x) a))
        ((<= (weight a) (weight (cadr x)))
         (generate (cdr x)
                   (huffman-node (car x) a)))
        (t
         (generate (insert a (cddr x))
                   (huffman-node (car x) (cadr x))))))

(defun generate-huffman-tree (x)
  (declare (xargs :guard (and (symbol-weight-pair-listp x)
                              (< 1 (len x)))
                  :verify-guards nil))
  (let ((leaves (pairs-to-leaves x)))
    (generate (cdr leaves) (car leaves))))

(defthm huffman-tree-listp-insert
  (implies (and (huffman-tree-listp x)
                (huffman-treep e))
           (huffman-tree-listp (insert e x))))

(defthm orderdp-insert
  (implies (orderdp x)
           (orderdp (insert e x))))

(defthm huffman-tree-listp-pairs-to-leaves
  (implies (symbol-weight-pair-listp x)
           (huffman-tree-listp (pairs-to-leaves x))))

(defthm orderdp-pairs-to-leaves
  (orderdp (pairs-to-leaves x)))

(defthm huffman-treep-generate
  (implies (and (huffman-treep a)
                (huffman-tree-listp x)
                (consp x))
           (huffman-treep (generate x a))))

(defthm nodep-generate
  (implies (and (huffman-treep a)
                (huffman-tree-listp x)
                (consp x))
           (nodep (generate x a))))

(defthm huffman-tree-listp-cdr
  (implies (huffman-tree-listp x)
           (huffman-tree-listp (cdr x))))

(defthm huffman-treep-car
  (implies (and (consp x) (huffman-tree-listp x))
           (huffman-treep (car x))))

(defthm len>1-consp
  (equal (< 1 (len x))
         (and (consp x)
              (consp (cdr x)))))

(defthm nodep-generate-huffman-tree
  (implies (and (symbol-weight-pair-listp x)
                (< 1 (len x)))
           (nodep (generate-huffman-tree x))))

(defthm huffman-treep-generate-huffman-tree
  (implies (and (symbol-weight-pair-listp x)
                (< 1 (len x)))
           (huffman-treep (generate-huffman-tree x))))

(verify-guards generate-huffman-tree)

(defun count-and-remove (e x)
  (declare (xargs :guard (true-listp x)))
  (cond ((endp x) (mv 0 nil))
        ((equal e (car x))
         (mv-let (c rest)
             (count-and-remove e (cdr x))
           (mv (+ 1 c) rest)))
        (t
         (mv-let (c rest)
             (count-and-remove e (cdr x))
           (mv c (cons (car x) rest))))))


(defun message-to-pairs (x)
  (declare (xargs :guard (true-listp x)))
  (if (endp x)
      nil
      (mv-let (c rest)
          (count-and-remove (car x) x)
        (cons (list (car x) c)
              (message-to-pairs rest)))))

(defthm rationalp-0-count-and-remove
  (rationalp (mv-nth 0 (count-and-remove e x))))

(defthm symbol-weight-pair-listp-message-to-pairs
  (symbol-weight-pair-listp (message-to-pairs x)))

;;; Huffman encodeing and decoding.
(defun encode-1 (x tree)
  (declare (xargs :guard (huffman-treep tree)
                  :measure (node-count tree)))
  (if (nodep tree)
      (cond ((member-equal x (symbol-list (left tree)))
             (cons 0 (encode-1 x (left tree))))
            (t
             (cons 1 (encode-1 x (right tree)))))
      nil))

(defun encode (x tree)
  (declare (xargs :guard (and (true-listp x)
                              (nodep tree)
                              (huffman-treep tree))))
  (if (endp x)
      nil
      (append (encode-1 (car x) tree)
              (encode (cdr x) tree))))

(defun bit-listp (x)
  (declare (xargs :guard t))
  (if (atom x)
      (null x)
      (and (bitp (car x))
           (bit-listp (cdr x)))))

(defun decode-1 (x tree)
  (declare (xargs :guard (and (bit-listp x)
                              (huffman-treep tree))
                  :measure (node-count tree)))
  (cond ((and (consp x) (nodep tree))
         (if (equal (car x) 0)
             (decode-1 (cdr x) (left tree))
             (decode-1 (cdr x) (right tree))))
        ((nodep tree) (mv nil nil))
        (t
         (mv (leaf-symbol tree) x))))

(defun decode (x tree)
  (declare (xargs :measure (len x)
                  :guard (and (bit-listp x)
                              (nodep tree)
                              (huffman-treep tree))))
  (if (or (endp x)
          (not (nodep tree)))
      nil
      (mv-let (symbol rest)
          (decode-1 x tree)
        (cons symbol (decode rest tree)))))

(defthm bit-listp-encode-1
  (bit-listp (encode-1 x tree)))

(defthm bit-listp-encode
  (bit-listp (encode x tree)))

(defthm consp-append
  (equal (consp (append x y))
         (or (consp x)
             (consp y))))

(defthm consp-encode-1
  (implies (and (huffman-treep tree)
                (nodep tree))
           (consp (encode-1 x tree))))

(defthm consp-encode
  (implies (and (consp x)
                (true-listp x)
                (huffman-treep tree)
                (nodep tree))
           (consp (encode x tree))))


(defthm len-decode-1
  (implies (and (consp x)
                (nodep tree))
           (< (len (mv-nth 1 (decode-1 x tree)))
              (len x))))

(defthm member-append
  (iff (member-equal e (append x y))
       (or (member-equal e x)
           (member-equal e y))))

(defthm decode-1-encode-1
  (implies (and (huffman-treep tree)
                (member-equal x (symbol-list tree)))
           (and (equal (mv-nth 0 (decode-1 (encode-1 x tree) tree))
                       x)
                (equal (mv-nth 1 (decode-1 (encode-1 x tree) tree))
                       nil))))

(defthm decode1-append-encode-1
  (implies (and (huffman-treep tree)
                (member-equal x (symbol-list tree)))
           (and (equal (mv-nth 0 (decode-1 (append (encode-1 x tree)
                                                   y)
                                           tree))
                       x)
                (equal (mv-nth 1 (decode-1 (append (encode-1 x tree)
                                                   y)
                                           tree))
                       y))))

(defthm decode-encode-to-decode1-encode1
  (implies (and (huffman-treep tree)
                (nodep tree)
                (member-equal x (symbol-list tree))
                (subsetp y (symbol-list tree)))
           (equal (decode (append (encode-1 x tree)
                                  (encode y tree))
                          tree)
                  (cons (mv-nth 0 (decode-1 (encode-1 x tree) tree))
                        (decode (encode y tree)
                                tree))))
  :hints (("Goal" :induct (decode x tree))))

(defthm decode-encode
  (implies (and (true-listp x)
                (huffman-treep tree)
                (nodep tree)
                (subsetp x (symbol-list tree)))
           (equal (decode (encode x tree)
                          tree)
                  x)))