-- Copyright (C) 2011 Tyson Brown
-- 
-- This file is part QuestHelper2, an AddOn for World of Warcraft.
-- 
-- QuestHelper2 is free software: you can redistribute it and/or modify
-- it under the terms of the GNU General Public License as published by
-- the Free Software Foundation, either version 3 of the License, or
-- (at your option) any later version.
-- 
-- QuestHelper2 is distributed in the hope that it will be useful,
-- but WITHOUT ANY WARRANTY; without even the implied warranty of
-- MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
-- GNU General Public License for more details.
-- 
-- You should have received a copy of the GNU General Public License
-- along with QuestHelper2. If not, see <http://www.gnu.org/licenses/>.

if LibDebug then LibDebug() end
local shared = select(2, ...)

-- This is an implementation if a binary min-heap, and is intended to be a replacement
-- for most of the uses currently implemented by the functions in binarysearch.lua.
-- It allows fast access to the lowest value in a set, and is much faster than the
-- binarysearch functions if you don't need the entire set to be sorted at all times.
-- 
-- This heap implementation also maintains a lookup table for each of the elements in
-- the heap, so that arbitrary elements can be removed or updated.
-- 
-- Unlike the binarysearch functions, these heap functions aren't stable -
-- the order of equivalent objects is undefined.

local floor = math.floor

local function defaultCmp(a, b)
  return a < b
end

function shared.createHeap(cmp)
  cmp = cmp or defaultCmp
  
  local heap, lookup, size = {}, {}, 0
  
  --[===[@debug@
  -- This sanity test is very slow.
  local function sanity()
    for i=1,size do
      assert(lookup[heap[i]] == i, "lookup error")
    end
    
    assert(heap[size+1] == nil, "extra elements")
    
    for obj, index in pairs(lookup) do
      assert(index >= 1 and index <= size, "stale lookup")
    end
    
    for i=2,size do
      assert(cmp(heap[floor(i/2)], heap[i]), "ranking error")
    end
  end
  --@end-debug@]===]
  
  -- Moves a value toward the root of the tree if necessary; and updates lookup for that variable.
  local function demote(i)
    --[===[@debug@
    assert(i >= 1 and i <= size)
    --@end-debug@]===]
    
    local temp = heap[i]
    
    while i ~= 1 do
      local prev = floor(i/2)
      local prevv = heap[prev]
      
      if cmp(prevv, temp) then
        lookup[temp] = i
        return
      end
      
      lookup[prevv], i, heap[i], heap[prev] = i, prev, prevv, temp
    end
    
    lookup[temp] = i
  end
  
  -- Moves a value away from the root of the tree if necessary; and updates lookup for that variable.
  local function promote(i)
    --[===[@debug@
    assert(i >= 1 and i <= size)
    --@end-debug@]===]
    
    local temp = heap[i]
    
    repeat
      local next = i*2
      local next2 = next+1
      local nextv, nextv2 = heap[next], heap[next2]
      
      if nextv and cmp(nextv, temp) then
        if nextv2 and cmp(nextv2, nextv) then
          lookup[nextv2], i, heap[i], heap[next2] = i, next2, nextv2, temp
        else
          lookup[nextv], i, heap[i], heap[next] = i, next, nextv, temp
        end
      else
        if nextv2 and cmp(nextv2, temp) then
          lookup[nextv2], i, heap[i], heap[next2] = i, next2, nextv2, temp
        else
          lookup[temp] = i
          return
        end
      end
    until false
  end
  
  -- Removes the next element from the heap and returns it.
  -- Returns nil if the heap is empty.
  function heap.pop()
    local value = heap[1]
    
    if size < 2 then
      if size == 1 then
        size, lookup[value], heap[1] = 0, nil, nil
      end
      
      return value
    end
    
    size, heap[1], heap[size], lookup[value] = size-1, heap[size], nil, nil
    promote(1)
    
    --[===[@debug@
    sanity()
    --@end-debug@]===]
    
    return value
  end
  
  -- Returns the next element in the heap, or nil if the heap is empty.
  function heap.next()
    return heap[1]
  end
  
  -- Returns true if the heap contains no elements.
  function heap.empty()
    return size == 0
  end
  
  -- Returns the index of a value in the heap, or nil if the heap does not contain that value.
  function heap.find(value)
    return lookup[value]
  end
  
  -- Inserts a value into the heap. You must not insert nil, false, or multiple instances of the same value.
  -- You may insert equivalent values, however the order they will be retrieved in will be undefined.
  function heap.insert(value)
    --[===[@debug@
    assert(value) -- We can't store nil or false; sorry.
    assert(lookup[value] == nil, "duplicate value: "..tostring(value))
    --@end-debug
    
    size = size+1
    heap[size] = value
    demote(size)
    
    --@debug@
    sanity()
    --@end-debug@]===]
  end
  
  -- Removes a value from the heap. The value must exist in the heap.
  function heap.remove(value)
    local i = lookup[value]
    
    --[===[@debug@
    assert(i)
    --@end-debug@]===]
    
    if i == size then
      size, heap[i], lookup[value] = size-1, nil, nil
    else
      heap[i], size, heap[size], lookup[value] = heap[size], size-1, nil, nil
      promote(i)
    end
    
    --[===[@debug@
    sanity()
    --@end-debug@]===]
  end
  
  -- If you modify a value in the heap such that its ranking relative to the other values in
  -- the heap may have changed, you must call this function on that value.
  -- It is logistically equivalent to removing that value and then re-inserting it.
  function heap.update(value)
    local i = lookup[value]
    
    --[===[@debug@
    assert(i)
    --@end-debug@]===]
    
    if i > 1 then 
      local prev = floor(i/2)
      local prevv = heap[prev]
      
      if cmp(value, prevv) then
        lookup[prevv], heap[i], heap[prev] = i, prevv, value
        demote(prev)
        --[===[@debug@
        sanity()
        --@end-debug@]===]
        return
      end
    end
    
    promote(i)
    
    --[===[@debug@
    sanity()
    --@end-debug@]===]
  end
  
  -- Treats the heap as an unsorted list, and swaps the elements so as to make a valid heap again.
  -- It is logistically equivalent to removing every element of the heap and then re-inserting them.
  function heap.heapify()
    size = #heap
    
    for i=2,size do
      demote(i)
    end
    
    --[===[@debug@
    sanity()
    --@end-debug@]===]
  end
  
  -- This function removes all values from the heap.
  function heap.clear()
    for value, index in pairs(lookup) do
      heap[index], lookup[value] = nil, nil
    end
    
    size = 0
  end
  
  return heap
end
