-- Copyright (C) 2010 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, ...)

local functions = {}

shared.meta.QuadTree = {
  __index = functions,
  name = "QuadTree"
}

function shared.createQuadTree()
  
  return setmetatable({info={}, dirty=true}, shared.meta.QuadTree)
end

local function rectCollision(x1, y1, x2, y2, x3, y3, x4, y4)
  return x1 <= x4 and x3 <= x2 and y1 <= y4 and y3 <= y2
end

local function addObject(self, object, x1, y1, x2, y2, x3, y3, x4, y4, info, depth)
  local xmid, ymid = (x3+x4)*0.5, (y3+y4)*0.5
  
  local subregions = self.a
  
  if not self.a then
    -- This region doesn't have any sub-regions.
    
    local objects = self.objects
    
    if not objects then
      -- It also doesn't contain any objects. Don't worry, we'll fix that.
      self.objects = {[object]=true}
      self.object_count = 1
      return
    end
    
    assert(not objects[object])
    objects[object] = true
    
    local count = self.object_count + 1
    
    if depth > 8 or count <= 8 then
      -- We're either too deeply nested with our subregions to create more,
      -- or we don't contain enough stuff to necessitate creating more subregions,
      -- to hold it all, so return.
      self.object_count = count
      return
    end
    
    -- Let's subdivide this region and give our objects to them.
    local a, b, c, d = {}, {}, {}, {}
    depth, self.a, self.b, self.c, self.d = depth+1, a, b, c, d
    
    for obj in pairs(objects) do
      local p = info[obj]
      
      x1, y1, x2, y2 = p.x1, p.y1, p.x2, p.y2
      
      if rectCollision(x1, y1, x2, y2,   x3,   y3,   xmid, ymid) then addObject(a, obj, x1, y1, x2, y2, x3,   y3,   xmid, ymid, info, depth) end
      if rectCollision(x1, y1, x2, y2,   xmid, y3,   x4,   ymid) then addObject(b, obj, x1, y1, x2, y2, xmid, y3,   x4,   ymid, info, depth) end
      if rectCollision(x1, y1, x2, y2,   x3,   ymid, xmid, y4  ) then addObject(c, obj, x1, y1, x2, y2, x3,   ymid, xmid, y4  , info, depth) end
      if rectCollision(x1, y1, x2, y2,   xmid, ymid, x4,   y4  ) then addObject(d, obj, x1, y1, x2, y2, xmid, ymid, x4,   y4  , info, depth) end
    end
    
    self.objects = nil
    self.object_count = nil
  else
    -- Add the object to the appropriate subregion.
    depth = depth + 1
    if rectCollision(x1, y1, x2, y2,   x3,   y3,   xmid, ymid) then addObject(self.a, object, x1, y1, x2, y2, x3,   y3,   xmid, ymid, info, depth) end
    if rectCollision(x1, y1, x2, y2,   xmid, y3,   x4,   ymid) then addObject(self.b, object, x1, y1, x2, y2, xmid, y3,   x4,   ymid, info, depth) end
    if rectCollision(x1, y1, x2, y2,   x3,   ymid, xmid, y4  ) then addObject(self.c, object, x1, y1, x2, y2, x3,   ymid, xmid, y4  , info, depth) end
    if rectCollision(x1, y1, x2, y2,   xmid, ymid, x4,   y4  ) then addObject(self.d, object, x1, y1, x2, y2, xmid, ymid, x4,   y4  , info, depth) end
  end
end

local function removeObject(self, object, x1, y1, x2, y2, x3, y3, x4, y4)
  if self.a then
    -- We have subregions; check if any of them intersect the object and ask them to remove the object if it does.
    local xmid, ymid, a, b, c, d = (x3+x4)*0.5, (y3+y4)*0.5, self.a, self.b, self.c, self.d
    
    if rectCollision(x1, y1, x2, y2,   x3,   y3,   xmid, ymid) then removeObject(a, object, x1, y1, x2, y2, x3,   y3,   xmid, ymid) end
    if rectCollision(x1, y1, x2, y2,   xmid, y3,   x4,   ymid) then removeObject(b, object, x1, y1, x2, y2, xmid, y3,   x4,   ymid) end
    if rectCollision(x1, y1, x2, y2,   x3,   ymid, xmid, y4  ) then removeObject(c, object, x1, y1, x2, y2, x3,   ymid, xmid, y4  ) end
    if rectCollision(x1, y1, x2, y2,   xmid, ymid, x4,   y4  ) then removeObject(d, object, x1, y1, x2, y2, xmid, ymid, x4,   y4  ) end
    
    if not (a.a or b.a or c.a or d.a) then
      -- None of our subregions have subregions, let's see if we could contain all their objects safely.
      
      -- NOTE: This doesn't account for the subregions containing the same object because it overlaps multiple of them;
      -- so the count may be less that what we calculate here, causing us a chance to merge them.
      -- Personally, I'm fine with that.
      local count = (a.object_count or 0) + (b.object_count or 0) + (c.object_count or 0) + (d.object_count or 0)
      
      if count <= 8 then
        -- We can hold all those objects, therefore we don't need our subregions.
        local objects, count = {}, 0
        self.objects = objects
        
        -- Hierarchy cat is in your directed graph, stealing your leaf nodes.
        if a.objects then for obj in pairs(a.objects) do if not objects[obj] then object[obj] = true count = count + 1 end end end
        if b.objects then for obj in pairs(b.objects) do if not objects[obj] then object[obj] = true count = count + 1 end end end
        if c.objects then for obj in pairs(c.objects) do if not objects[obj] then object[obj] = true count = count + 1 end end end
        if d.objects then for obj in pairs(d.objects) do if not objects[obj] then object[obj] = true count = count + 1 end end end
        
        self.object_count, self.a, self.b, self.c, self.d = count, nil, nil, nil, nil
      end
    end
    
    return
  end
  
  local objects = self.objects
  if objects then
    -- Check if we have the object in question.
    if objects[object] then
      objects[object] = nil
      local count = self.object_count
      if count == 1 then self.objects, self.object_count = nil, 0
      else self.object_count = count - 1 end
    end
  end
end

local function iterateObjects(self, x1, y1, x2, y2, x3, y3, x4, y4, info, set)
  if self.a then
    -- Check the sub-regions.
    local xmid, ymid = (x3+x4)*0.5, (y3+y4)*0.5
    if rectCollision(x1, y1, x2, y2,   x3,   y3,   xmid, ymid) then iterateObjects(self.a, x1, y1, x2, y2,   x3,   y3,   xmid, ymid, info, set) end
    if rectCollision(x1, y1, x2, y2,   xmid, y3,   x4,   ymid) then iterateObjects(self.b, x1, y1, x2, y2,   xmid, y3,   x4,   ymid, info, set) end
    if rectCollision(x1, y1, x2, y2,   x3,   ymid, xmid, y4  ) then iterateObjects(self.c, x1, y1, x2, y2,   x3,   ymid, xmid, y4  , info, set) end
    if rectCollision(x1, y1, x2, y2,   xmid, ymid, x4,   y4  ) then iterateObjects(self.d, x1, y1, x2, y2,   xmid, ymid, x4,   y4  , info, set) end
    return
  end
  
  -- Check if we have any objects that intersect the given rectangle.
  local objects = self.objects
  if objects then for obj in pairs(objects) do
    -- Only do the collision test if another region hasn't already added it.
    if not set[obj] then
      local p = info[obj]
      if rectCollision(x1, y1, x2, y2, p.x1, p.y1, p.x2, p.y2) then set[obj] = true end
    end
  end end
end

function functions:addObject(object,x1,y1,x2,y2)
  if not x2 then
    -- Adding a point; emulate it using a rectangle with zero area.
    x2,y2 = x1,y1
  end
  
  local info = self.info
  local old_info = info[object]
  
  if old_info then
    -- We're really moving the object.
    if not self.dirty then
      if old_info.x1 == self.x1 or old_info.y1 == self.y1 or old_info.x2 == self.x2 or old_info.y2 == self.y2 then
        -- This object was touching our bounding box, it may need to be recalculated.
        self.dirty = true
      else
        removeObject(self, object, old_info.x1, old_info.y1, old_info.x2, old_info.y2, self.x1, self.y1, self.x2, self.y2)
      end
    end
    
    old_info.x1, old_info.y1, old_info.x2, old_info.y2 = x1, y1, x2, y2
  else
    -- We're adding the object, for reals.
    self.info[object] = {x1=x1, y1=y1, x2=x2, y2=y2}
  end
  
  if not self.dirty then
    if x1 < self.x1 or y1 < self.x1 or x2 > self.x2 or y2 > self.y2 then
      -- This is outside our bounding box, it's going to need to be recalculated.
      self.dirty = true
    else
      addObject(self, object, x1, y1, x2, y2, self.x1, self.y1, self.x2, self.y2, info, 1)
    end
  end
end

function functions:removeObject(object)
  local old_info = self.info[object]
  
  if old_info then
    self.info[object] = nil
    
    if not self.dirty then
      if old_info.x1 == self.x1 or old_info.y1 == self.y1 or old_info.x2 == self.x2 or old_info.y2 == self.y2 then
        -- This object was touching our bounding box, it may need to be recalculated.
        self.dirty = true
      else
        removeObject(self, object, old_info.x1, old_info.y1, old_info.x2, old_info.y2, self.x1, self.y1, self.x2, self.y2)
      end
    end
  end
end

function functions:clear()
  table.wipe(self.info)
  self.dirty = true
end

function functions:boundingRectangle(object)
  local info = self.info[object]
  
  if info then
    return info.x1, info.y1, info.x2, info.y2
  end
end

local function nop()
end

function functions:iterateObjects(x1,y1,x2,y2)
  local info = self.info
  
  if self.dirty then
    if not next(info) then
      -- We contain no objects, the bounding rectangle is undefined.
      -- Return an iterator that returns no objects rather than
      -- trying to rebuild the tree.
      return nop
    end
    
    -- Calculate a bounding rectangle that will cover all the objects we contain.
    local min, max = math.min, math.max
    local x1, y1, x2, y2
    
    for object, rect in pairs(info) do
      if not x1 then
        x1, y1, x2, y2 = rect.x1, rect.y1, rect.x2, rect.y2
      else
        x1, y1, x2, y2 = min(x1, rect.x1), min(y1, rect.y1), max(x2, rect.x2), max(y2, rect.y2)
      end
    end
    
    self.x1, self.y1, self.x2, self.y2 = x1, y1, x2, y2
    
    -- Clear any old data that we don't need anymore.
    self.a, self.b, self.c, self.d, self.objects, self.object_count = nil, nil, nil, nil, nil, nil
    
    -- Regenerate the quad tree.
    for object, rect in pairs(info) do
      addObject(self, object, rect.x1, rect.y1, rect.x2, rect.y2, x1, y1, x2, y2, info, 1)
    end
    
    -- We are now clean. Horray!
    self.dirty = nil
  end
  
  if not x2 then
    -- Adding a point; emulate it using a rectangle with zero area.
    x2,y2 = x1,y1
  end
  
  local set = {}
  
  iterateObjects(self, x1, y1, x2, y2, self.x1, self.y1, self.x2, self.y2, self.info, set)
  
  local object = next(set)
  
  return function ()
    local result = object
    if object then object = next(set, object) end
    return result
  end
end
