/-
Copyright (c) 2019 Microsoft Corporation. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Leonardo de Moura
-/
prelude
import init.data.list init.coe
universes u

namespace BinomialHeapImp

structure HeapNodeAux (α : Type u) (h : Type u) :=
(val : α) (rank : Nat) (children : List h)

inductive Heap (α : Type u) : Type u
| empty {} : Heap
| heap  (ns : List (HeapNodeAux α Heap)) : Heap

abbrev HeapNode (α) := HeapNodeAux α (Heap α)

variables {α : Type u}

instance : Inhabited (Heap α) := ⟨Heap.empty⟩

def hRank : List (HeapNode α) → Nat
| []   => 0
| h::_ => h.rank

def isEmpty : Heap α → Bool
| Heap.empty => true
| _          => false

def singleton (a : α) : Heap α :=
Heap.heap [{ val := a, rank := 1, children := [] }]

@[specialize] def combine (lt : α → α → Bool) (n₁ n₂ : HeapNode α) : HeapNode α :=
if lt n₂.val n₁.val then
   { rank := n₂.rank + 1, children := n₂.children ++ [Heap.heap [n₁]], .. n₂ }
else
   { rank := n₁.rank + 1, children := n₁.children ++ [Heap.heap [n₂]], .. n₁ }

@[specialize] partial def mergeNodes (lt : α → α → Bool) : List (HeapNode α) → List (HeapNode α) → List (HeapNode α)
| [], h  => h
| h,  [] => h
| f@(h₁ :: t₁), s@(h₂ :: t₂) =>
  if h₁.rank < h₂.rank then h₁ :: mergeNodes t₁ s
  else if h₂.rank < h₁.rank then h₂ :: mergeNodes t₂ f
  else
    let merged := combine lt h₁ h₂;
    let r      := merged.rank;
    if r != hRank t₁ then
      if r != hRank t₂ then merged :: mergeNodes t₁ t₂ else mergeNodes (merged :: t₁) t₂
    else
      if r != hRank t₂ then mergeNodes t₁ (merged :: t₂) else merged :: mergeNodes t₁ t₂

@[specialize] def merge (lt : α → α → Bool) : Heap α → Heap α → Heap α
| Heap.empty,    h           => h
| h,             Heap.empty  => h
| Heap.heap h₁, Heap.heap h₂ => Heap.heap (mergeNodes lt h₁ h₂)

@[specialize] def headOpt (lt : α → α → Bool) : Heap α → Option α
| Heap.empty  => none
| Heap.heap h => h.foldl
  (fun r n => match r with
   | none   => n.val
   | some v => if lt v n.val then v else n.val) none

/- O(log n) -/
@[specialize] def head [Inhabited α] (lt : α → α → Bool) : Heap α → α
| Heap.empty        => default α
| Heap.heap []      => default α
| Heap.heap (h::hs) => hs.foldl (fun r n => if lt r n.val then r else n.val) h.val

@[specialize] def findMin (lt : α → α → Bool) : List (HeapNode α) → Nat → HeapNode α × Nat → HeapNode α × Nat
| [],    _,   r          => r
| h::hs, idx, (h', idx') => if lt h.val h'.val then findMin hs (idx+1) (h, idx) else findMin hs (idx+1) (h', idx')

def tail (lt : α → α → Bool) : Heap α → Heap α
| Heap.empty        => Heap.empty
| Heap.heap []      => Heap.empty
| Heap.heap [h]     =>
  match h.children with
  | []      => Heap.empty
  | (h::hs) => hs.foldl (merge lt) h
| Heap.heap hhs@(h::hs) =>
  let (min, minIdx) := findMin lt hs 1 (h, 0);
  let rest          := hhs.eraseIdx minIdx;
  min.children.foldl (merge lt) (Heap.heap rest)

partial def toList (lt : α → α → Bool) : Heap α → List α
| Heap.empty => []
| h          => match headOpt lt h with
  | none   => []
  | some a => a :: toList (tail lt h)

inductive WellFormed (lt : α → α → Bool) : Heap α → Prop
| emptyWff                  : WellFormed Heap.empty
| singletonWff (a : α)      : WellFormed (singleton a)
| mergeWff (h₁ h₂ : Heap α) : WellFormed h₁ → WellFormed h₂ → WellFormed (merge lt h₁ h₂)
| tailWff (h : Heap α)      : WellFormed h → WellFormed (tail lt h)

end BinomialHeapImp

open BinomialHeapImp

def BinomialHeap (α : Type u) (lt : α → α → Bool) := { h : Heap α // WellFormed lt h }

@[inline] def mkBinomialHeap (α : Type u) (lt : α → α → Bool) : BinomialHeap α lt :=
⟨Heap.empty, WellFormed.emptyWff lt⟩

namespace BinomialHeap
variables {α : Type u} {lt : α → α → Bool}

@[inline] def empty : BinomialHeap α lt :=
mkBinomialHeap α lt

@[inline] def isEmpty : BinomialHeap α lt → Bool
| ⟨b, _⟩ => BinomialHeapImp.isEmpty b

/- O(1) -/
@[inline] def singleton (a : α) : BinomialHeap α lt :=
⟨BinomialHeapImp.singleton a, WellFormed.singletonWff lt a⟩

/- O(log n) -/
@[inline] def merge : BinomialHeap α lt → BinomialHeap α lt → BinomialHeap α lt
| ⟨b₁, h₁⟩, ⟨b₂, h₂⟩ => ⟨BinomialHeapImp.merge lt b₁ b₂, WellFormed.mergeWff b₁ b₂ h₁ h₂⟩

/- O(log n) -/
@[inline] def head [Inhabited α] : BinomialHeap α lt → α
| ⟨b, _⟩ => BinomialHeapImp.head lt b

/- O(log n) -/
@[inline] def headOpt : BinomialHeap α lt → Option α
| ⟨b, _⟩ => BinomialHeapImp.headOpt lt b

/- O(log n) -/
@[inline] def tail : BinomialHeap α lt → BinomialHeap α lt
| ⟨b, h⟩ => ⟨BinomialHeapImp.tail lt b, WellFormed.tailWff b h⟩

/- O(log n) -/
@[inline] def insert (a : α) (h : BinomialHeap α lt) : BinomialHeap α lt :=
merge (singleton a) h

/- O(n log n) -/
@[inline] def toList : BinomialHeap α lt → List α
| ⟨b, _⟩ => BinomialHeapImp.toList lt b

end BinomialHeap