rj-action-library/Runtime/Structures/Spatial/QuadTree/QuadTree.cs

using System.Collections;
using System.Collections.Generic;
using Godot;
using System;
using System.Linq;



namespace Rokojori
{
  public class QuadTree<T>:QuadTreeNode<T>
  { 
    protected Func<T,Vector2> _getPosition;
    protected Func<List<T>,List<T>> _combinePoints;
    protected Func<List<T>,List<T>, float,List<T>> _smoothPoints;

    protected Vector2 _min;
    protected Vector2 _max;
    protected float _rootCellSize;

    protected int _maxDepth = 64;
    public int maxDepth => _maxDepth;

    public QuadTree( Func<T,Vector2> getPosition, Vector2 min, Vector2 max, float rootCellSize, int maxDepth  )
    {
      this._getPosition = getPosition;
      this._min = min;
      this._max = max;
      this._rootCellSize = rootCellSize;
      this._maxDepth = maxDepth;
      
      CreateRootCells();

    }

     public QuadTree( Vector2 center, float rootCellSize, int maxDepth  )
    {
      this._min = center - Vector2.One * rootCellSize / 2f;
      this._max = center + Vector2.One * rootCellSize / 2f;
      this._rootCellSize = rootCellSize;
      this._maxDepth = maxDepth;
      
      CreateRootCells();

    }

    public void SetCombiner( Func<List<T>,List<T>> combinePoints )
    {
      this._combinePoints = combinePoints;
    }

    public void SetSmoother( Func<List<T>,List<T>,float,List<T>> smoothPoints )
    {
      this._smoothPoints = smoothPoints;
    }


    public List<T> CombinePoints( List<T> values )
    {
      return _combinePoints( values );
    }

    public List<T> SmoothPoints( List<T> targetValues, List<T> sourceSmoothing, float amount )
    {
      return _smoothPoints( targetValues, sourceSmoothing, amount );
    }




    List<QuadTreeCell<T>> _rootCells = new List<QuadTreeCell<T>>();
    public List<QuadTreeCell<T>> rootCells => _rootCells;

    public bool Insert( List<T> data )
    {
      var result = true;

      data.ForEach( 
        ( d )=>
        {
          result = result && Insert( d );
        }
      );

      return result;
    }

    public List<T> GetInsideBox( Box2 box, List<T> list = null )
    {
      list = list == null ? new List<T>() : list;

      GetInsideBoxWithDuplicates( box, list );

      var set = new HashSet<T>( [ .. list ] );
      return set.ToList<T>();
    } 

    public List<T> GetInsideBoxWithDuplicates( Box2 box, List<T> list = null )
    {
      list = list == null ? new List<T>() : list;
      var rootIndexMin = PositionToRootIndex( box.min );
      var rootIndexMax = PositionToRootIndex( box.max );

      for ( int x = rootIndexMin.X; x <= rootIndexMax.X; x++ )
      {
        for ( int y = rootIndexMin.Y; y <= rootIndexMax.Y; y++ )
        {
          
          var rootIndex = new Vector2I( x, y );
          var rootCell = GetRootCellByRootIndex( rootIndex );

          rootCell.GetInsideBox( box, list );
          
        }
      }

      return list;

    }

    public bool Insert( T data )
    {
      var rootCell = GetRootCell( _getPosition( data ) );

      if ( rootCell == null )
      {
        return false;
      }

      if ( ! rootCell.box.ContainsPoint( GetPosition( data ) ) )
      {
        RJLog.Log( "Box not containing point", rootCell.rootCellIndex );
      }

      return rootCell.Insert( data );
    }

    public void BoxInsert( T data, Box2 box )
    {
      var rootIndexMin = PositionToRootIndex( box.min );
      var rootIndexMax = PositionToRootIndex( box.max );

      for ( int x = rootIndexMin.X; x <= rootIndexMax.X; x++ )
      {
        for ( int y = rootIndexMin.Y; y <= rootIndexMax.Y; y++ )
        {
          
          var rootIndex = new Vector2I( x, y );
          var rootCell = GetRootCellByRootIndex( rootIndex );

          rootCell.BoxInsert( data, box );
        
        }
      }
    }

    public int GetLevelForSize( float size )
    {
      var level = 0;
      var it = _rootCellSize;

      while ( it > size )
      {
        level ++;
        it /= 2f;
      }

      return level;
    }

    public float GetSizeForLevel( int level )
    {
      var it = 0;
      var size = _rootCellSize;
    
      while ( it < level )
      {
        size /= 2f;
        it ++;
      }

      return size;
    }

    public void CombineUntilSize( float size )
    {
      CombineUntilLevel( GetLevelForSize( size ) );
    } 
  

    public void CombineUntilLevel( int minLevel, float smoothDown = 0.2f )
    { 
      var levelMap = GetLevelMap();
      var levels = levelMap.Keys.ToList();
      levels.Sort();
      var maxLevel = levels.Last();

      for ( int i = maxLevel - 1; i > minLevel; i-- )
      {
        levelMap[ i ].ForEach(
          ( c )=>
          {
            c.Combine();
          }
        );
      }     

      if ( smoothDown > 0 )
      {
        for ( int i = minLevel + 1; i < maxLevel - 1; i++ )
        {
          levelMap[ i ].ForEach(
            ( c )=>
            {
              c.SmoothDown( smoothDown );
            }
          );
        }
      }
    }

    public MapList<int,QuadTreeCell<T>> GetLevelMap()
    {
      var walker = new QuadTreeWalker<T>();

      var levelMap = new MapList<int,QuadTreeCell<T>>();

      walker.Iterate(
        this, 
        ( n )=>
        {
          if ( n is QuadTree<T> tree )
          {
            return;
          }

          var c = n as QuadTreeCell<T>;

          levelMap.Add( c.depth, c );
        }
      );    

      return levelMap;
    }

    Vector2 _startOffset = Vector2.Zero;
    Vector2I _rootCellDimensions = Vector2I.Zero;

    public Vector2I PositionToRootIndex( Vector2 position )
    {
      position -= _startOffset;

      var index = ( position / ( Vector2.One * _rootCellSize ) ).FloorToInt();

      return index;
    }

    public Vector2 GetPosition( T data )
    {
      return _getPosition( data );
    }

    public QuadTreeCell<T> GetRootCell( Vector2 position )
    {
      var rootIndex = PositionToRootIndex( position );
      return GetRootCellByRootIndex( rootIndex );
    }
    
    public QuadTreeCell<T> GetRootCellByRootIndex( Vector2I rootCellIndex )
    {
      var rootCellFlatIndex = MathX.MultiIndexToFlatIndex( rootCellIndex, _rootCellDimensions );

      if ( rootCellFlatIndex < 0 || rootCellFlatIndex >= _rootCells.Count )
      {
        return null;
      }

      return _rootCells[ rootCellFlatIndex ];
    }
    

    void CreateRootCells()
    {
      var rootCellSize3 = Vector2.One * _rootCellSize;
      var start = _min.SnapFloored( rootCellSize3 );
      var end =   _max.SnapCeiled( rootCellSize3 );

      var num = ( ( end - start ) / rootCellSize3 ).RoundToInt();

      _rootCellDimensions = num;
      _startOffset = start;

      for ( int x = 0; x < num.X ; x ++ )
      {
        for ( int y = 0; y < num.Y; y++ )
        {         
          var rootIndex = new Vector2I( x, y );
          var min = rootIndex * rootCellSize3 + start;
          var max = min + rootCellSize3;

          var cell = QuadTreeCell<T>.Create( this, min, max, _rootCells.Count );

          _rootCells.Add( cell );            

        }
      }
    }
    
  }
}