#pragma once

#include <vector>
#include <cstdint>
#include <J3ML/LinearAlgebra/Vector2.h>
#include "AABB2D.h"

namespace Geometry {
    using LinearAlgebra::Vector2;

    template<typename T>
    class QuadTree {
    public:
        struct Node {
            Node *parent;
            uint32_t childIndex;
            std::vector<T> objects;

            Vector2 center;
            Vector2 radius;

            bool IsLeaf() const { return childIndex == 0xFFFFFFFF; }

            uint32_t TopLeftChildIndex() const { return childIndex; }

            uint32_t TopRightChildIndex() const { return childIndex + 1; }

            uint32_t BottomLeftChildIndex() const { return childIndex + 2; }

            uint32_t BottomRightChildIndex() const { return childIndex + 3; }

            /// This assumes that the QuadTree contains unique objects and never duplicates
            void Remove(const T &object) {
                for (size_t i = 0; i < objects.size(); ++i) {
                    if (objects[i] == object) {
                        AssociateQuadTreeNode(object,
                                              0); // Mark in the object that it has been removed from the QuadTree.
                        std::swap(objects[i], objects.back());
                        objects.pop_back();
                        return;
                    }
                }
            }

            AABB2D ComputeAABB() {}

            float DistanceSq(const Vector2 &point) const {
                Vector2 centered = point - center;
                Vector2 closestPoint = centered.Clamp(-radius, radius);
                return closestPoint.DistanceSq(centered);
            }

        };

        // Helper struct used when traversing through the tree
        struct TraversalStackItem {
            Node *node;
        };

        QuadTree() :
                rootNodeIndex(-1),
                boundingAABB(Vector2(0, 0), Vector2(0, 0)) {
            // TODO: currently storing persistent raw pointers to this array outside the array
            nodes.reserve(200000);
        }

        /// Removes all nodes and objects in this tree and reintializes the tree to a single root node.
        void Clear(const Vector2 &minXY = Vector2(-1, -1), const Vector2 &maxXY = Vector2(1, 1));

        /// Places the given object onto the proper (leaf) node of the tree. After placing, if the leaf split rule is
        /// satisfied, subdivides the leaf node into 4 subquadrants and reassings the objects to the new leaves.
        void Add(const T &object);

        /// Removes the given object from this tree.
        /// To call this function, you must define a function QuadTree<T>::Node *GetQuadTreeNode(const T& object)
        /// which returns the node of this quadtree where the object resides in.
        void Remove(const T &object);

        /// @return The bounding rectangle for the whole tree.
        /// @note This bounding rectangle does not tightly bound the objects themselves, only the root node of the tree.
        AABB2D BoundingAABB() const { return boundingAABB; }

        /// @return The topmost node in the tree.
        Node *Root();

        const Node *Root() const;

        /// Returns the total number of nodes (all nodes, i.e. inner nodes + leaves) in the tree.
        /// Runs in constant time.
        int NumNodes() const;

        /// Returns the total number of leaf nodes in the tree.
        /// @warning Runs in time linear 'O(n)' to the number of nodes in the tree.
        int NumLeaves() const;

        /// Returns the total number of inner nodes in the tree.
        /// @warning Runs in time linear 'O(n)' to the number of nodes in the tree.
        int NumInnerNodes() const;

        /// Returns the total number of objects stored in the tree.
        /// @warning Runs in time linear 'O(n)' to the number of nodes in the tree.
        int NumObjects() const;

        /// Returns the maximum height of the whole tree (the path from the root to the farthest leaf node).
        int TreeHeight() const;

        /// Returns the height of the subtree rooted at 'node'.
        int TreeHeight(const Node *node) const;

        /// Performs an AABB intersection query in this Quadtreee, and calls the given callback function for each non-empty
        /// node of the tree which intersects the given AABB.
        /** @param aabb The axis-aligned bounding box to intersect this QuadTree with.
            @param callback A function or a function object of prototype
                bool callbackFunction(QuadTree<T> &tree, const AABB2D &queryAABB, QuadTree<T>::Node &node);
            If the callback function returns true, the execution of the query is stopped and this function immediately
            returns afterwards. If the callback function returns false, the execution of the query continues. */
        template<typename Func>
        inline void AABBQuery(const AABB2D &aabb, Func &callback);

        /// Finds all object pairs inside the given AABB which have colliding AABBs. For each such pair, calls the
        /// specified callback function.
        template<typename Func>
        inline void CollidingPairsQuery(const AABB2D &aabb, Func &callback);

        /// Performs various consistency checks on the given node. Use only for debugging purposes.
        void DebugSanityCheckNode(Node *n);

    private:
        void Add(const T &object, Node *n);

        /// Allocates a sequential 4-tuple of QuadtreeNodes, contiguous in memory.
        int AllocateNodeGroup(Node *parent);

        void SplitLeaf(Node *leaf);

        std::vector<Node> nodes;
        int rootNodeIndex;
        AABB2D boundingAABB;

        void GrowRootTopLeft();

        void GrowRootTopRight();

        void GrowRootBottomLeft();

        void GrowRootBottomRight();

        void GrowImpl(int quadrantForRoot);
    };

    // NOTE: Keep members defined here. Template-parameterized classes
    // can't be split across header and implementation files
    // but the presence of the implementation file is a requirement


    template<typename T>
    void QuadTree<T>::Clear(const Vector2 &minXY, const Vector2 &maxXY) {
        nodes.clear();

        boundingAABB.minPoint = minXY;
        boundingAABB.maxPoint = maxXY;

        rootNodeIndex = AllocateNodeGroup(0);

        Node *root = Root();
        root->center = (minXY + maxXY) * 0.5f;
        root->radius = maxXY - root->center;
    }

    template<typename T>
    void QuadTree<T>::Add(const T &object) {
        Node *n = Root();
        AABB2D objectAABB = GetAABB2D(object);

        // Ramen Noodle Bowls of nested if statements are generally bad practice
        // Unfortunately, sometimes the problem domain makes this unavoidable

        if (objectAABB.minPoint.x >= boundingAABB.minPoint.x) {
            // object fits left.
            if (objectAABB.maxPoint.x <= boundingAABB.maxPoint.x) {
                // object fits left and right.
                if (objectAABB.minPoint.y >= boundingAABB.minPoint.y) {
                    // Object fits left, right, and top.
                    if (objectAABB.maxPoint.y <= boundingAABB.maxPoint.y) {
                        // Object fits the whole root AABB. Can safely add into the existing tree size.
                        AddObject(object, n);
                        return;
                    } else {
                        // Object fits left, right, and top, but not bottom.
                        GrowRootBottomRight(); // Could grow bottom-left as well, wouldn't matter here.
                    }
                } else {
                    // Object fits left and right, but not to top.
                    GrowRootTopRight(); // Could grow top-left as well, wouldn't matter here.
                }
            } else {
                // Object fits left, but not to right. We must grow right. Check whether to grow top or bottom;
                if (objectAABB.minPoint.y < boundingAABB.minPoint.y)
                    GrowRootTopRight();
                else
                    GrowRootBottomRight();
            }
        } else {
            // We must grow left. Check whether to grow top or bottom.
            if (objectAABB.minPoint.y < boundingAABB.minPoint.y)
                GrowRootTopLeft();
            else
                GrowRootBottomLeft();
        }
        // Now that we have grown the tree root node, try adding again.
        Add(object);
    }

    template<typename T>
    void QuadTree<T>::Remove(const T &object) {
        Node *n = GetQuadTreeNode(object);
        if (n) {
            n->Remove(object);
        }
    }

    template<typename T>
    void QuadTree<T>::Add(const T &object, Node *n) {
        for (;;) {
            // Traverse the QuadTree to decide which quad to place this object on.
            float left = n->center.x - MinX(object); // If left > 0.f, then the object overlaps with the left quadrant.
            float right = MaxX(object) - n->center.x; // If right > 0.f, then the object overlaps with the right quadrant.

            float top = n->center.y - MinY(object); // If top > 0.f, then the object overlaps with the top quadrant
            float bottom =
                    MaxY(object) - n->center.y; // If bottom > 0.f, then the object overlaps with the bottom quadrant
            float leftAndRight = std::min(left, right); // If > 0.f, then the object straddles left-right halves.
            float topAndBottom = std::min(top, bottom); // If > 0.f, then the object straddles top-bottom halves.
            float straddledEitherOne = std::max(leftAndRight,
                                                topAndBottom); // If > 0.f, thne the object is in two or more quadrants.

            // Note: It can happen that !left && !right, or !top && !bottom.
            // but the if()s are setup below so that right/bottom is taken if no left/top, so that is ok.

            // We must put the object onto this node if
            // a) the object straddled the parent->child split lines.
            // b) this object is a leaf
            if (straddledEitherOne > 0.f) {
                n->objects.push_back(object);
                AssociateQuadTreeNode(object, n);
                return;
            }
            if (n->IsLeaf()) {
                n->objects.push_back(object);
                AssociateQuadTreeNode(object, n);

                if ((int) n->objects.size() > minQuadTreeNodeObjectCount &&
                    std::min(n->radius.x, n->radius.y) >= minQuadTreeQuadrantSize) {
                    SplitLeaf(n);
                }

                return;
            }
            if (left > 0.f) {
                if (top > 0.f) {
                    n = &nodes[n->TopLeftChildIndex()];
                } else {
                    n = &nodes[n->BottomLeftChildIndex()];
                }
            } else {
                if (top > 0.f) {
                    n = &nodes[n->TopRightChildIndex()];
                } else {
                    n = &nodes[n->BottomRightChildIndex()];
                }
            }
        }
    }

    template<typename T>
    typename QuadTree<T>::Node *QuadTree<T>::Root() {
        return nodes.empty() ? 0 : &nodes[rootNodeIndex];
    }

    template<typename T>
    const typename QuadTree<T>::Node *QuadTree<T>::Root() const {
        return nodes.empty() ? 0 : &nodes[rootNodeIndex];
    }

    template <typename T>
    int QuadTree<T>::AllocateNodeGroup(Node* parent) {
        size_t oldCap = nodes.capacity();

        int index = (int)nodes.size();
        Node n;
        n.parent = parent;
        n.childIndex = 0xFFFFFFFF;
        if (parent) {
            n.radius = parent->radius * 0.5f;
            n.center = parent->center - n.radius;
        }

        nodes.push_back(n);
        if (parent)
            n.center.x = parent->center.x + n.radius.x;
        nodes.push_back(n);

        if (parent) {
            n.center.x = parent->center.x - n.radius.x;
            n.center.y = parent->center.y + n.radius.y;
        }

        nodes.push_back(n);
        if (parent)
            n.center.x = parent->center.x + n.radius.x;
        nodes.push_back(n);

        return index;
    }

    template <typename T>
    void QuadTree<T>::SplitLeaf(Node *leaf) {
        leaf->childIndex = AllocateNodeGroup(leaf);

        size_t i = 0;
        while (i < leaf->objects.size()) {
            const T& object = leaf->objects[i];

            // Traverse the QuadTree to decide which quad to place this object into
            float left = leaf->center.x - MinX(object);
            float right = MaxX(object) - leaf->center;
            float top = leaf->center.y - MinY(object);
            float bottom = MaxY(object) - leaf->center.y;
            float leftAndRight = std::min(left, right);
            float topAndBottom = std::min(top, bottom);

            float straddledEitherOne = std::max(leftAndRight, topAndBottom);

            if (straddledEitherOne > 0.f) {
                ++i;
                continue;
            }

            if (left > 0.f) {
                if (top > 0.f) {
                    Add(object, &nodes[leaf->TopLeftChildIndex()]);
                } else {
                    Add(object, &nodes[leaf->BottomLeftChildIndex()]);
                }
            } else {
                if (top > 0.f) {
                    Add(object, &nodes[leaf->TopRightChildIndex()]);
                } else {
                    Add(object, &nodes[leaf->BottomRightChildIndex()]);
                }
            }

            // Remove the object we added to a child from this node.
            leaf->objects[i] = leaf->objects.back();
            leaf->objects.pop_back();
        }
    }
}