#include <vector>
#include <cmath>
#include <cstdio>

/*

 TODO: probably should generalize this to N-dimensions  (quadtree is 2D specific case - octree is 3D specific case)
 
 Idea - could quadtree be used to optimize a 2D UI system?

 Multiple ideas to test:
   - able to register callbacks if nodes touched or events happen to them
   - be able to detect collisions by when adding an object to a node, can check the other objects already in the node
   - fast optimizations for moving an object from one node to another - eg: add then remove, or descend to common node of
         the add/removes, and then do the add and removes from the common node down
   - having multiple trees for different object types
         - different ideas here - one is to have completely different trees
         - another idea is to have one tree but each node contains a set of bit fields which tells which types of objects are below
         - trade-off between searching for specific item types or for searching for anything in a given region
   - having bit-fields in the nodes for some kind of parameter of the contents of the nodes
         - eg: 64-bits per node for flags, the flags could be one bit per object type
   - LINQ system - having a system to test predicates - speed up searches
   - resouvoir sampling for getting some representative set of samples
   - query a region - circle / sphere - CSG shapes - handle objects having dimensions - eg: add a 3x3x3 sized object, and be able
        to find it when searching a given region

 Applications:
    - trigger volumes
    - collisions
    - height-map for positioning to ground
    - finding objects in an area
    - finding pick-ups
    - frustum clipping

*/

template <typename T>
class Flags
{
public:
  Flags(T val) : value((int)val) {}
  operator T()
  {
    return (T)value;
  }
  inline Flags<T>& operator|=(T val)
  {
    value |= (int)val;
    return *this;
  }
  int value;
};


class Quadtree;
class QuadtreeItemInterface
{
public:
  QuadtreeItemInterface(Quadtree* tree, float qX, float qY)
  {
    // add to quad tree
  }

  virtual ~QuadtreeItemInterface()
  {
    // needs to be removed from the quadtree
  }

  // need an interface for moving it

  virtual float quadtreeX() = 0;
  virtual float quadtreeY() = 0;
  //virtual void NodeAccessed() = 0;

  Quadtree* tree;
  int qX, qY; // qt pos stored at - perhaps could use morton code
};


class Quadtree
{
public:
  Quadtree(float resolution = 1.0f);
  ~Quadtree();

  // probably needs to be a shared_ptr or may be a way to remove it from the quadtree when it is destroyed
  void AddItem(QuadtreeItemInterface* item)
  {
    auto addFunctor = [](QuadtreeNode& node, QuadtreeItemInterface* item) { node.AddItem(item); };
    QuadtreePoint pnt = QuantizePoint(item->quadtreeX(), item->quadtreeY());
    Expand(pnt);
    Traverse(addFunctor, *rootNode, bounds, pnt, item);
  }

  void RemoveItem(QuadtreeItemInterface* item) // need to quickly test the item belongs to the quadtree
  {
    // TODO: this doesn't recursively remove nodes as it traverses back up the tree
    //auto removeFunctor = [](QuadtreeNode& node, QuadtreeItemInterface* item) { node.RemoveItem(item); };
    //Traverse(removeFunctor, *rootNode, bounds, QuantizePoint(item->quadtreeX(), item->quadtreeY()), item);
    
    RemoveTraverse(*rootNode, bounds, QuantizePoint(item->quadtreeX(), item->quadtreeY()), item);
  }

  void MoveItem(QuadtreeItemInterface* item, float x, float y)
  {
    QuadtreePoint oldPnt = QuantizePoint(item->quadtreeX(), item->quadtreeY());
    QuadtreePoint newPnt = QuantizePoint(x, y);
    MoveTraverse(*rootNode, bounds, oldPnt, newPnt, item);
  }

private:
  enum class Direction : int 
  {
    Middle = 0,
    Left = 1<<0,
    Right = 1<<1,
    Down = 1<<2,
    Up = 1<<4,
    UpLeft = Up | Left,
    UpRight = Up | Right,
    DownLeft = Down | Left,
    DownRight = Down | Right
  };

  struct QuadtreePoint
  {
    int x, y;
  };

  struct QuadtreeBounds
  {
    QuadtreePoint origin;
    int size;
  };

  class QuadtreeNode
  {
  public:
    QuadtreeNode* nodes[4];
    void AddItem(QuadtreeItemInterface* item)
    {
      items.push_back(item);
    }
    void RemoveItem(QuadtreeItemInterface* item)
    {
      items.erase(std::remove(items.begin(), items.end(), item), items.end());
    }
    size_t ItemCount()
    {
      return items.size();
    }
  private:
    std::vector<QuadtreeItemInterface*> items;
  };

  float resolution;
  QuadtreeNode* rootNode = nullptr;
  QuadtreeBounds bounds;

  int QuantizeValue(float value)
  {
    return (int)std::floor((value - resolution * 0.5f) / resolution);
  }

  QuadtreePoint QuantizePoint(float x, float y)
  {
    return (QuadtreePoint){ QuantizeValue(x), QuantizeValue(y) };
  }

  void Expand(const QuadtreePoint& pnt)
  {
    if (!rootNode)
    {
      rootNode = new QuadtreeNode;
      bounds.origin = pnt;
      bounds.size = 1;
      return;
    }
    ExpandRecurse(pnt);
  }

  static inline QuadtreeBounds ChildBounds(const QuadtreeBounds& b, int quadrant)
  {
    QuadtreeBounds childBounds = b;
    childBounds.size = b.size >> 1;
    childBounds.origin.x += (quadrant & 1) ? childBounds.size : 0;
    childBounds.origin.y += (quadrant & 2) ? 0 : childBounds.size;
    return childBounds;
  }
  
  static inline bool PointInBounds(const QuadtreePoint& pnt, const QuadtreeBounds& b)
  {
    return (pnt.x >= b.origin.x && pnt.y >= b.origin.y && pnt.x < (b.origin.x + b.size) && pnt.y < (b.origin.y + b.size));
  }

  bool RemoveTraverse(QuadtreeNode& node, const QuadtreeBounds& b, const QuadtreePoint& pnt, QuadtreeItemInterface* item)
  {
    if (b.size == 1)
    {
      node.RemoveItem(item);
      return node.ItemCount() == 0;
    }
    for (int q = 0; q < 4; q++)
    {
      QuadtreeBounds childBounds = ChildBounds(b, q);
      if (PointInBounds(pnt, childBounds))
      {
        //assert(node.nodes[q]);
        if (RemoveTraverse(*node.nodes[q], childBounds, pnt, item))
        {
          delete node.nodes[q];
          node.nodes[q] = 0;
        }
        break;
      }
    }
    return !node.nodes[0] && !node.nodes[1] && !node.nodes[2] && !node.nodes[3];
  }


  void MoveTraverse(QuadtreeNode& node, const QuadtreeBounds& b, const QuadtreePoint& oldPnt, const QuadtreePoint& newPnt, QuadtreeItemInterface* item)
  {
    //assert(b.size != 1)
    for (int q = 0; q < 4; q++)
    {
      QuadtreeBounds childBounds = ChildBounds(b, q);
      bool oldIn = PointInBounds(oldPnt, childBounds);
      bool newIn = PointInBounds(newPnt, childBounds);
      if (oldIn && newIn)
      {
        //assert(node.nodes[q]);
        MoveTraverse(*node.nodes[q], childBounds, oldPnt, newPnt, item);
      }
      else if (oldIn)
      {
        //assert(node.nodes[q]);
        RemoveTraverse(*node.nodes[q], childBounds, oldPnt, item);
      }
      else if (newIn)
      {
        auto addFunctor = [](QuadtreeNode& node, QuadtreeItemInterface* item) { node.AddItem(item); };
        Traverse(addFunctor, *node.nodes[q], childBounds, newPnt, item);
      }
    }
  }


  template <typename F>
  void Traverse(const F& func, QuadtreeNode& node, const QuadtreeBounds& b, const QuadtreePoint& pnt, QuadtreeItemInterface* item)
  {
    if (b.size == 1)
    {
      func(node, item);
      return;
    }
    for (int q = 0; q < 4; q++)
    {
      QuadtreeBounds childBounds = ChildBounds(b, q);
      if (PointInBounds(pnt, childBounds))
      {
        if (!node.nodes[q])
          node.nodes[q] = new QuadtreeNode;
        Traverse(func, *node.nodes[q], childBounds, pnt, item);
        return;
      }
    }

    /*
    // Generic traverse idea - callbacks for each part of the traversal
    auto actionFunctor = [func, item](QuadtreeNode& node, const QuadtreeBounds& bounds) { func(node, item); };
    auto terminateFunctor = [](const QuadtreeNode& node, const QuadtreeBounds& bounds) { return bounds.size == 1; };
    auto testFunctor = [pnt](const QuadtreeNode& node, const QuadtreeBounds& bounds) { return PointInBounds(pnt, bounds); };
    auto nullNodeFunctor = [](QuadtreeNode& node, int quadrant) { node.nodes[quadrant] = new QuadtreeNode; };
    TraverseGeneric(actionFunctor, terminateFunctor, testFunctor, nullNodeFunctor, node, b, true);
    */
  }

  template <typename ActionF, typename TerminateF, typename TestF, typename NullNode>
  void TraverseGeneric(const ActionF& action, const TerminateF& term, const TestF& test, const NullNode& nullNode, QuadtreeNode& node, const QuadtreeBounds& b, bool earlyOut)
  {
    if (term(node, b))
    {
      action(node, b);
      return;
    }
    for (int q = 0; q < 4; q++)
    {
      QuadtreeBounds childBounds = ChildBounds(b, q);
      if (test(node, childBounds))
      {
        if (!node.nodes[q])
          nullNode(node, q);
        if (!node.nodes[q])
          TraverseGeneric(action, term, test, nullNode, *node.nodes[q], childBounds, earlyOut);
        if (earlyOut)
          return;
      }
    }
  }

  void ExpandRecurse(const QuadtreePoint& pnt)
  {
    // assert(rootNode);
    Flags<Direction> growDirectionFlags = Direction::Middle;
    growDirectionFlags |= (pnt.x < bounds.origin.x) ? Direction::Left : Direction::Middle;
    growDirectionFlags |= (pnt.x >= (bounds.origin.x + bounds.size)) ? Direction::Right : Direction::Middle;
    growDirectionFlags |= (pnt.y < bounds.origin.y) ? Direction::Down : Direction::Middle;
    growDirectionFlags |= (pnt.y >= (bounds.origin.y + bounds.size)) ? Direction::Up : Direction::Middle;
    int oldRootQuadrant = 0;
    switch (growDirectionFlags) {
      case Direction::Middle:    return; // Nothing more to do, the point is already inside the current bounds
      case Direction::Up:
      case Direction::UpLeft:    oldRootQuadrant = 3; break;
      case Direction::Right:
      case Direction::UpRight:   oldRootQuadrant = 2; break;
      case Direction::Left:
      case Direction::DownLeft:  oldRootQuadrant = 1; break;
      case Direction::Down:
      case Direction::DownRight: oldRootQuadrant = 0; break;
    }
    int q = oldRootQuadrant;
    bounds.origin.x -= (q & 1) ? bounds.size : 0;
    bounds.origin.y -= (q & 2) ? 0 : bounds.size;
    bounds.size <<= 1;
    QuadtreeNode* newRootNode = new QuadtreeNode;
    newRootNode->nodes[q] = rootNode;
    rootNode = newRootNode;
    ExpandRecurse(pnt);
  }
};