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

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



namespace Rokojori
{
  public class OcTree<T,D>:OcTreeNode<T,D>
  { 
    protected Func<T,Vector3> _getPosition;
    protected Func<List<T>,List<T>> _combinePoints;
    protected Func<List<T>,List<T>, float,List<T>> _smoothPoints;

    protected Vector3 _min;
    protected Vector3 _max;
    protected float _rootCellSize;

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

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

    }

     public OcTree( Vector3 center, float rootCellSize, int maxDepth  )
    {
      this._min = center - Vector3.One * rootCellSize / 2f;
      this._max = center + Vector3.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<OcTreeCell<T,D>> _rootCells = new List<OcTreeCell<T,D>>();
    public List<OcTreeCell<T,D>> rootCells => _rootCells;

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

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

      return result;
    }

    public List<T> GetInsideBox( Box3 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( Box3 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++ )
        {
          for ( int z = rootIndexMin.Z; z <= rootIndexMax.Z; z++ )
          {
            var rootIndex = new Vector3I( x, y, z );
            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, Box3 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++ )
        {
          for ( int z = rootIndexMin.Z; z <= rootIndexMax.Z; z++ )
          {
            var rootIndex = new Vector3I( x, y, z );
            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,OcTreeCell<T,D>> GetLevelMap()
    {
      var walker = new OcTreeWalker<T,D>();

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

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

          var c = n as OcTreeCell<T,D>;

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

      return levelMap;
    }

    Vector3 _startOffset = Vector3.Zero;
    Vector3I _rootCellDimensions = Vector3I.Zero;

    public Vector3I PositionToRootIndex( Vector3 position )
    {
      position -= _startOffset;

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

      return index;
    }

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

    public OcTreeCell<T,D> GetRootCell( Vector3 position )
    {
      var rootIndex = PositionToRootIndex( position );
      return GetRootCellByRootIndex( rootIndex );
    }
    
    public OcTreeCell<T,D> GetRootCellByRootIndex( Vector3I rootCellIndex )
    {
      var rootCellFlatIndex = MathX.MultiIndexToFlatIndex( rootCellIndex, _rootCellDimensions );

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

      return _rootCells[ rootCellFlatIndex ];
    }
    

    void CreateRootCells()
    {
      var rootCellSize3 = Vector3.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++ )
        {
          for ( int z = 0; z < num.Z; z++ )
          {
            var rootIndex = new Vector3I( x, y, z );
            var min = rootIndex * rootCellSize3 + start;
            var max = min + rootCellSize3;

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

            _rootCells.Add( cell );            

            // var positionRootIndex = PositionToRootIndex( cell.center );
            
            // RJLog.Log( 
            //   "Creating root cell", 
            //   "Roots Index:", _rootCells.Count - 1, 
            //   "Cell Index:", cell.rootCellIndex, 
            //   "Center:", cell.center,
            //   "RootIndex:", rootIndex,
            //   "Get Self ByRootIndex:", GetRootCellByRootIndex( rootIndex ) == cell,
            //   "RootIndex From Position: ", positionRootIndex,
            //   "Get Self ByRootIndex From Position:", GetRootCellByRootIndex( positionRootIndex ) == cell
            // );
          }
        }
      }
    }
    
  }
}