Firebird_2D/Firebird2D.Datatypes/Quadtree.cs

using Firebird2D.Requirement;
using SFML.Graphics;
using System.Diagnostics.CodeAnalysis;
using Firebird2D.SFMLExtensions;
using System.Drawing;
using SFML.System;

namespace Firebird2D.Datatypes
{
    public class Quadtree<T> where T : ISpatial
    {
        private readonly List<T> _elements = [];

        private readonly int _bucketCapacity;

        private readonly int _maxDepth;

        private Quadtree<T>? _topLeft, _topRight, _bottomLeft, _bottomRight;

        public FloatRect Bounds { get; }

        public int Level { get; init; }

        [MemberNotNullWhen(false, nameof(_topLeft), nameof(_topRight), nameof(_bottomLeft), nameof(_bottomRight))]
        public bool IsLeaf
            => _topLeft == null || _topRight == null || _bottomLeft == null || _bottomRight == null;

        public bool AllowOutOfBounds { get; set; }

        public Quadtree(FloatRect bounds, int bucketCapacity, int maxDepth)
        {
            Bounds = bounds;
            _bucketCapacity = bucketCapacity;
            _maxDepth = maxDepth;
        }

        public Quadtree(FloatRect bounds): this(bounds,32,5) { }

        public void Insert(T element)
        {
            Require.NotNull(element, nameof(element));

            if (!(Bounds.Contains(element.Bounds) || AllowOutOfBounds))
                throw new ArgumentException(string.Format("{0} is out of Quadtreebounds", nameof(element)), nameof(element));

            if (_elements.Count >= _bucketCapacity)
                Split();

            Quadtree<T>? containingChild = GetContainingChild(element.Bounds);

            if (containingChild != null) 
            {
                containingChild.Insert(element);
            }
            else
            {
                _elements.Add(element);
            }
        }

        public bool Remove(T element)
        {
            Require.NotNull(element, nameof(element));

            Quadtree<T>? containingChild = GetContainingChild(element.Bounds);

            bool removed = containingChild?.Remove(element) ?? _elements.Remove(element);

            if (removed && CountElements() <= _bucketCapacity)
                Merge();

            return removed;
        }

        public int CountElements()
        {
            int count = _elements.Count;
            if(!IsLeaf)
            {
                count += _topLeft.CountElements();
                count += _topRight.CountElements();
                count += _bottomLeft.CountElements();
                count += _bottomRight.CountElements();
            }

            return count;
        }

        public IEnumerable<T> GetElements() 
        { 
            List<T> children = new();
            Queue<Quadtree<T>> nodes = new Queue<Quadtree<T>>();

            nodes.Enqueue(this);

            while (nodes.Count > 0)
            {
                Quadtree<T> node = nodes.Dequeue();

                if (!node.IsLeaf)
                {
                    nodes.Enqueue(node._topLeft);
                    nodes.Enqueue(node._topRight);
                    nodes.Enqueue(node._bottomLeft);
                    nodes.Enqueue(node._bottomRight);
                }

                children.AddRange(node._elements);
            }

            return children;
        }

        private void Split()
        {
            if (!IsLeaf)
                return;

            if (Level + 1 > _maxDepth)
                return;

            _topLeft = CreateChild(Bounds.Location());
            _topRight = CreateChild(new Vector2f(Bounds.Center().X, Bounds.Location().Y));
            _bottomLeft = CreateChild(new Vector2f(Bounds.Location().X, Bounds.Center().Y));
            _bottomRight = CreateChild(Bounds.Center());

            List<T> elements = _elements.ToList();

            foreach (T element in elements)
            {
                Quadtree<T>? containingChild = GetContainingChild(element.Bounds);

                if (containingChild != null)
                {
                    _elements.Remove(element);

                    containingChild.Insert(element);
                }
            }
        }

        private Quadtree<T> CreateChild(Vector2f location)
            => new(new FloatRect(location, Bounds.Size / 2), _bucketCapacity, _maxDepth) { Level = Level + 1 };

        private void Merge()
        {
            if (IsLeaf)
                return;

            _elements.AddRange(_topLeft._elements);
            _elements.AddRange(_topRight._elements);
            _elements.AddRange(_bottomLeft._elements);
            _elements.AddRange(_bottomRight._elements);

            _topLeft = _topRight = _bottomLeft = _bottomRight = null;
        }

        private Quadtree<T>? GetContainingChild(FloatRect bounds)
        {
            if (IsLeaf)
                return null;

            if(_topLeft.Bounds.Contains(bounds))
                return _topLeft;

            if (_topRight.Bounds.Contains(bounds))
                return _topRight;

            if (_bottomLeft.Bounds.Contains(bounds))
                return _bottomLeft;

            if (_bottomRight.Bounds.Contains(bounds))
                return _bottomRight;

            return null;
        }

    }
}