; Abstract heap with user-supplied ordering function
; an ordering function is a binary relation
; `<` will create a MinHeap, `>` will create a MaxHeap
;
; for any two items `a` and `b`,
;   `a` will be higher in the Heap then `b` if `(ord a b)`
(doc Heap "is a heap data structure that uses a user-supplied ordering function
to order its values.")
(defmodule Heap
    (hidden lchild)
    (defn lchild [i]
        (+ 1 (* 2 i)))

    (hidden rchild)
    (defn rchild [i]
        (+ 2 (* 2 i)))

    (hidden parent)
    (defn parent [i]
        (/ (- i 1) 2))

    (hidden max-of-three-until!)
    (doc max-of-three-until! "Get the index for the largest (by ord) of an element and its two children.")
    (defn max-of-three-until! [heap i len ord]
        (let-do [lchild-i (lchild i)
                 rchild-i (rchild i)]
                 (when (and (< lchild-i len) (~ord (Array.unsafe-nth heap lchild-i) (Array.unsafe-nth heap i)))
                       (set! i lchild-i))
                 (when (and (< rchild-i len) (~ord (Array.unsafe-nth heap rchild-i) (Array.unsafe-nth heap i)))
                       (set! i rchild-i))
                 i))

    ; push-down-until!, push-down!, and push-up! are intended for use only by
    ; those performing internal mutation to the heap who want to restore order.

    (defn push-down-until! [heap i len ord]
        (while true
               (let [challenger (max-of-three-until! heap i len ord)]
                    (if (= challenger i)
                        (break)
                        (do
                            (Array.swap! heap i challenger)
                            (set! i challenger))))))

    (defn push-down! [heap i ord]
        (push-down-until! heap i (Array.length heap) ord))

    (defn push-up! [heap i ord]
        (while (/= i 0)
            (let [elem (Array.unsafe-nth heap i)
                  parent-i (Heap.parent i)
                  parent-elem (Array.unsafe-nth heap parent-i)]
                 (if (not (~ord elem parent-elem))
                     (break)
                     (do (Array.swap! heap i parent-i)
                         (set! i parent-i))))))

    (doc peek "Returns first item on heap.")
    (defn peek [heap]
        (Array.first heap))

    (doc heapify! "Convert array to a heap in place")
    (defn heapify! [arr ord]
        (let [len (Array.length arr)]
            (for [i 1 len]
                (push-up! arr i ord))))

    (doc push! "Insert a new item onto the heap.")
    (defn push! [heap item ord]
        (do
            (Array.push-back! heap item)
            (push-up! heap (- (Array.length heap) 1) ord)))

    (doc pop! "Remove and return the first item in the heap.")
    (defn pop! [heap ord]
        (do
            ; swap 0 with tail
            (Array.swap! heap 0 (- (Array.length heap) 1))
            ; restore heap excluding tail
            (push-down-until! heap 0 (- (Array.length heap) 1) ord)
            ; pop off tail, returning it
            (Array.pop-back! heap)))
)

(doc MinHeap "is a heap that uses `<` to order its values.")
(defmodule MinHeap
    (hidden ord)
    (defn ord [a b]
        (< a b))

    (defn push-down! [heap i]
        (Heap.push-down! heap i &ord))

    (defn push-down-until! [heap i len]
        (Heap.push-down-until! heap i len &ord))

    (defn push-up! [heap i]
        (Heap.push-up! heap i &ord))

    (doc peek "Returns minimum item on min-heap.")
    (defn peek [heap]
        (Array.first heap))

    (doc heapify! "Convert array to a min-heap in place")
    (defn heapify! [arr]
        (Heap.heapify! arr &ord))

    (doc push! "Insert a new element onto the min-heap.")
    (defn push! [heap item]
        (Heap.push! heap item &ord))

    (doc pop! "Remove and return the first item in the min-heap.")
    (defn pop! [heap]
        (Heap.pop! heap &ord))
)

(doc MaxHeap "is a heap that uses `>` to order its values.")
(defmodule MaxHeap
    (hidden ord)
    (defn ord [a b]
        (> a b))

    (defn push-down! [heap i]
        (Heap.push-down! heap i &ord))

    (defn push-down-until! [heap i len]
        (Heap.push-down-until! heap i len &ord))

    (defn push-up! [heap i]
        (Heap.push-up! heap i &ord))

    (doc peek "Returns maximum first item on max-heap.")
    (defn peek [heap]
        (Array.first heap))

    (doc heapify! "Convert array to a max-heap in place")
    (defn heapify! [arr]
        (Heap.heapify! arr &ord))

    (doc push! "Insert a new element onto the max-heap.")
    (defn push! [heap item]
        (Heap.push! heap item &ord))

    (doc pop! "Remove and return the first item in the max-heap.")
    (defn pop! [heap]
        (Heap.pop! heap &ord))
)

(doc HeapSort "is a module for sorting using the heap data structure.")
(defmodule HeapSort
    (hidden ord)
    (defn ord [a b]
        (> a b))

    (doc sort-by! "Perform an in-place heapsort of a given array with a comparison function.")
    (defn sort-by! [arr f]
        (do
            (Heap.heapify! arr f)
            ; now we walk through the array, at all times 0..tail is a max heap
            ; and tail..len is the sorted output
            ; we slowly grow the tail while shrinking the head
            ; once tail == 0, then we have finished sorting
            (let [tail (- (Array.length arr) 1)]
                 (while (> tail 0)
                        (do
                            ; swap values of [0] and [tail]
                            ; this makes our heap temporarily invalid
                            (Array.swap! arr 0 tail)
                            ; keep pushing down 0 until heap property is satisfied
                            (Heap.push-down-until! arr 0 tail f)
                            ; grow our tail, shrinking our head
                            (set! tail (- tail 1)))))))

    (doc sorted-by "Perform a heapsort in a new copy of given array with comparison function.")
    (defn sorted-by [arr f]
        (let-do [narr (Array.copy arr)]
                (sort-by! &narr f)
                narr))

    (doc sort-by "Perform an in-place heapsort of a given owned array with comparison function.")
    (defn sort-by [arr f]
        (do
            (sort-by! &arr f)
            arr))

    (doc sort! "Perform an in-place heapsort of a given array.")
    (defn sort! [arr]
      (sort-by! arr &ord))

    (doc sorted "Perform a heapsort in a new copy of given array.")
    (defn sorted [arr]
      (sorted-by arr &ord))

    (doc sort "Perform an in-place heapsort of a given owned array.")
    (defn sort [arr]
      (sort-by arr &ord))
)