require "./heap.cr"

# Inspired by Python's heapdict
class HeapHash(K,V)
  private class Wrapper(K,V)
    property key : K
    property value : V
    property index : Int32

    def initialize(@key, @value, @index)
    end
  end

  private class InternalHeap(K, V) < Array(Wrapper(K,V))
    def swap(i,j)
      super(i,j)
      self[i].index = i
      self[j].index = j
    end
  end

  def initialize
    @hash = {} of K => Wrapper(K,V)
    @heap = InternalHeap(K,V).new
  end

  def delete(k)
    wrapper = @hash[k]
    @heap.bubble_up(wrapper.index) do |i, j|
      false
    end
    pop
  end

  def []=(k : K, v : V)
    if @hash.has_key? k
      delete k
    end
    wrapper = Wrapper(K,V).new(k, v, @heap.size)
    @hash[k] = wrapper
    @heap.heap_push(wrapper) do |i, j|
      i.value < j.value
    end
  end

  def [](k)
    wrapper = @hash[k]
    wrapper.value
  end

  def []?(k)
    return nil unless wrapper = @hash[k]?
    wrapper.value
  end

  def pop
    wrapper = @heap.heap_pop do |l, r|
      l.value < r.value
    end
    @hash.delete wrapper.key
    {wrapper.key, wrapper.value}
  end

  delegate size, empty?, to: @heap
end