#pragma once



namespace J3ML::Geometry
{

    enum CardinalAxis
    {
        AxisX = 0,
        AxisY,
        AxisZ,
        AxisNone
    };


    struct KdTreeNode
    {
        /// If this is an inner node, specifies along which axis this node is split.
        /// If this is a leaf, has the value AxisNone.
        unsigned splitAxis  : 2;
        /// If this is an inner node, specifies the index/offset to the child node pair.
        /// If this is a leaf, the value is undefined.
        unsigned childIndex : 30;
        union
        {
            /// If this is an inner node, specifies the position along the cardinal axis of the split.
            float splitPos;
            /// If this is a leaf, specifies the index/ofset to the object bucket for this leaf.
            /// If zero, this leaf does not have a bucket associated with it. (empty leaf)
            u32 bucketIndex;
        };

        /// If true, this leaf does not contain any objects.
        bool IsEmptyLeaf() const { assert(IsLeaf()); return bucketIndex == 0; }
        bool IsLeaf() const { return splitAxis == AxisNone; }
        int LeftChildIndex() const { return (int)childIndex; }
        int RightChildIndex() const { return (int) childIndex+1; }
        CardinalAxis SplitAxis() const { return (CardinalAxis) splitAxis;}
    };

    /// A KD-tree accelleration structure for static geometry.
    template <typename T>
    class KdTree
    {
    public:
        KdTree() {}
        //~KDTree() { /* TODO: FILL */}
        void AddObjects(const T *objects, int numObjects);
        void Build();
        void Clear();
        u32* Bucket(int bucketIndex);
        const u32* Bucket(int bucketIndex) const;
        T& Object(int objectIndex);
        const T& Object(int objectIndex) const;

        int NumObjects() const;

        int NumNodes() const;
        int NumLeaves() const;
        int NumInnerNodes() const;
        int TreeHeight() const;
        KdTreeNode* Root();
        const KdTreeNode* Root() const;
        bool IsPartOfThisTree(const KdTreeNode *node) const;
        //const AABB& BoundingAABB() const { return rootAABB;}

        /// Traverses a ray through this kD-tree, and calls the given leafCallback function for each leaf of the tree.
        /// Uses the "recursive B" method from Vlastimil Havran's thesis.
        template <typename Func>
        void RayQuery(const Ray& r, Func &leaftCallback);

        template <typename Func>
        void AABBQuery(const AABB& aabb, Func& leafCallback);
    private:
        static const int maxNodes = 256 * 1024;
        static const int maxTreeDepth = 30;

        std::vector<KdTreeNode> nodes;
        //std::vector<u8, Aligned

    };
}