module ParasitemiaCore.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)

[<Struct>]
type private Node<'k, 'v> =
    val mutable key : 'k
    val value : 'v
    new (k, v) = { key = k; value = v }
    override this.ToString () = sprintf "%O -> %O" this.key this.value

/// <summary>
/// An heap min or max depending of the given key comparer.
/// The goal is to have a set of data and be able to get the value associated with the min (or max) key value.
/// </summary>
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

    // Check the integrity of the heap, use for debug only.
    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 node = a.[0]
        a.[0] <- a.[a.Count - 1]
        a.RemoveAt (a.Count - 1)
        heapUp 0
        node.key, node.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 (Node (key, 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 ()