module Heap

open System.Collections.Generic

let inline private parent (i: int) : int = (i - 1) / 2
let inline private left (i: int) : int = 2 * (i + 1) - 1
let inline private right (i: int) : int = 2 * (i + 1)

// TODO: measure performance with a struct.
type private Node<'k, 'v> = { mutable key: 'k; value : 'v } with
    override this.ToString () = sprintf "%A -> %A" this.key this.value

type Heap<'k, 'v> (kComparer : IComparer<'k>) =
    let a = List<Node<'k, 'v>>()

    let rec heapUp (i: int) =
        let l, r = left i, right i

        // Is the left child greater than the parent?
        let mutable max = if l < a.Count && kComparer.Compare(a.[l].key, a.[i].key) > 0 then l else i

        // Is the right child greater than the parent and the left child?
        if r < a.Count && kComparer.Compare(a.[r].key, a.[max].key) > 0
        then
            max <- r

        // If a child is greater than the parent.
        if max <> i
        then
            let tmp = a.[i]
            a.[i] <- a.[max]
            a.[max] <- tmp
            heapUp max

    let rec checkIntegrity (i: int) : bool =
        let l, r = left i, right i

        let leftIntegrity =
            if l < a.Count
            then
                if kComparer.Compare(a.[l].key, a.[i].key) > 0
                then false
                else checkIntegrity l
            else
                true
        let rightIntegrity =
            if r < a.Count
            then
                if kComparer.Compare(a.[r].key, a.[i].key) > 0
                then false
                else checkIntegrity r
            else
                true
        leftIntegrity && rightIntegrity

    interface IEnumerable<'k * 'v> with
        member this.GetEnumerator () : IEnumerator<'k * 'v> =
            (seq { for e in a -> e.key, e.value }).GetEnumerator()

    interface System.Collections.IEnumerable with
        member this.GetEnumerator () : System.Collections.IEnumerator =
            (this :> IEnumerable<'k * 'v>).GetEnumerator() :> System.Collections.IEnumerator

    member this.Next () : 'k * 'v =
        let { key = key; value = value } = a.[0]
        a.[0] <- a.[a.Count - 1]
        a.RemoveAt(a.Count - 1)
        heapUp 0
        key, value

    member this.RemoveNext () =
        a.[0] <- a.[a.Count - 1]
        a.RemoveAt(a.Count - 1)
        heapUp 0

    member this.Add (key: 'k) (value: 'v) =
        a.Add({ key = key; value = value })

        let mutable i = a.Count - 1
        while i > 0 && kComparer.Compare(a.[parent i].key, a.[i].key) < 0 do
            let tmp = a.[parent i]
            a.[parent i] <- a.[i]
            a.[i] <- tmp
            i <- parent i

    member this.IsEmpty = a.Count = 0
    member this.Count = a.Count

    member this.Max : 'k * 'v =
        let max = a.[0]
        max.key, max.value

    member this.Clear () = a.Clear()