rj-action-library/Runtime/Procedural/Assets/Tree/TreeGenerator.cs

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



namespace Rokojori
{ 
  [Tool]
  [GlobalClass, Icon("res://addons/rokojori_action_library/Icons/Scatterer.svg") ]
  public partial class TreeGenerator:Node3D
  {
    [Export]
    public float rootRadius;
  
    [Export]
    public float rootHeight;

    [Export]
    public Curve branchPosition;

    [Export]
    public Curve levelToChildBranches;

    [Export]
    public Curve levelToSubdivisionNoise;

    [Export]
    public Curve subdivisionNoiseRange;

    [Export]
    public int maxLevel = 4;

   
    

    [Export]
    public Curve radiusScale;

    [Export]
    public Curve heightScale;

    [Export]
    public Curve siblingRotationVariance;

    [Export]
    public Curve rotationPerLevel;

    [Export]
    public Curve rotationVariation;

    [Export]
    public int radialSegments = 8;

    [Export]
    public int minRadialSegements = 3;

    [Export]
    public int maxRadialSegments = 64;

    [Export]
    public Curve radialSegementsMultiplier;

    [Export]
    public bool avoidClusters = true;

    [Export]
    public float minClusterDistance = 1;

    [Export]
    public float noise = 0.1f;

    [Export]
    public TubeGeometrySettings branchSettings;

    [Export]
    public TubeShape[] shapes;

    [Export]
    public Material branchMaterial;

    [Export]
    public Material leavesMaterial;

    [Export]
    public int leavesRows = 4;

    [Export]
    public int leavesColumns = 4;

    [Export]
    public bool hideStructureWhenMeshWasBuild = true;

     [Export]
    public int seed = 1998;

    [Export]
    public bool randomizeSeed = false;

    [ExportToolButton( "Generate Branches Structure" )]
    public Callable GenerateBranchesButton => Callable.From(
      ( )=>
      {
        GenerateBranches();
      }
    );

    [ExportToolButton( "Generate Branches Mesh" )]
    public Callable GenerateMeshButton => Callable.From(
      ( )=>
      {
        GenerateMesh();
      }
    );

    [ExportToolButton( "Generate Branches And Mesh" )]
    public Callable GenerateBranchesMeshButton => Callable.From(
      ( )=>
      {
        GenerateBranchesAndMesh();
      }
    );

    [Export]
    bool _generating = false;

    [Export]
    public TreeBranch[] branches = [];

    List<TreeBranch> _branches = [];

    SeedableRandomEngine random;


    [Export]
    public Node deselecter;

    async Task GenerateBranchesAndMesh()
    {
      try
      {
        if ( useDebugSingleMeshGeneration )
        {
          EditorInterface.Singleton.GetSelection().Clear();
          EditorInterface.Singleton.GetSelection().AddNode( deselecter );
          await this.RequestNextFrame();
        }

        await GenerateBranches();
        await GenerateMesh();
      }
      catch( System.Exception e )
      {
        this.LogError( e );
      }


      return;
    }

    [Export]
    public bool useDebugSingleMeshGeneration = false;

    async Task GenerateMesh()
    { 
      var time = Async.StartTimer();

      var mg = new MeshGeometry();

      var meshInstances = this.GetDirectChildren<MeshInstance3D>();
      meshInstances.ForEach( m => m.SelfDestroy() );

      if ( useDebugSingleMeshGeneration )
      {
        this.GetDirectChildren<MeshInstance3D>().ForEach( m => m.SelfDestroy() );
      }
      else
      {
        this.GetAll<MeshInstance3D>().ForEach( m => m.SelfDestroy() );
      }

      for ( int i = 0; i < _branches.Count; i++ )
      {
        time = await Async.WaitIfExceeded( time, this, processTimePerFrame );
        var branchMG = CreateBranchMesh( _branches[ i ] );

        if ( useDebugSingleMeshGeneration )
        {
          var meshInstance3D = this.CreateChild<MeshInstance3D>( "Branch Mesh @" + _branches[ i ].index );
          meshInstance3D.Mesh = branchMG.GenerateMesh();
          meshInstance3D.SetSurfaceOverrideMaterial( 0, branchMaterial );

          if ( i % 5 == 0 )
          {
            await this.RequestNextFrame();
          }
          

        }
        else
        {
          mg.Add( branchMG );
        } 
        
      } 

      var leavesMG = new MeshGeometry();

      for ( int bb =0; bb < _branches.Count; bb++ )
      {
        time = await Async.WaitIfExceeded( time, this, processTimePerFrame );
        var b = _branches[ bb ];

        if ( b.level < leavesLevel )
        {
          continue;
        }
        

        var numLeaves = ( b.height / leafSize ) * random.Sample( leavesDensity );
        var curve = b.GetSpline( this ).GetCurve();

      

        for ( int i = 0; i < numLeaves; i++ )
        {
          var t = i / ( numLeaves - 1f );
          var leafGeometry = new MeshGeometry();

          var index = random.IntegerExclusive( leavesRows * leavesColumns );
          var indexX = index % leavesColumns;
          var indexY = index / leavesColumns;

          var uvStart = new Vector2( indexX, indexY ) / new Vector2( leavesColumns, leavesRows );
          var uvEnd = uvStart +  Vector2.One / new Vector2( leavesColumns, leavesRows );

          
          var p = curve.PoseAt( t );
          var rx = random.Range( 0, 360 );
          var ry = random.Range( 0, 360 );
          var rz = random.Range( 0, 360 );

          var leafSizeCurrent = leafSize * random.Sample( leafSizeScaleVariance, 1, 1 );

          leafGeometry.AddQuad( Math3D.FromEulerDegrees( new Vector3( rx, ry, rz ) ),
            leafSizeCurrent, uvStart, uvEnd, p.position );

          leafGeometry.BlendNormalsOverY( Vector3.Up, leavesNormalBlending, leavesStartBlend, leavesEndBlend, leavesBlendCurve ); 

          if ( useDebugSingleMeshGeneration )
          {
            time = await Async.WaitIfExceeded( time, processTimePerFrame );
             var leavesInstance = this.CreateChild<MeshInstance3D>( "Leaves Mesh @" + b.index + "." + i );
            leavesInstance.Mesh = leafGeometry.GenerateMesh();
            leavesInstance.SetSurfaceOverrideMaterial( 0, leavesMaterial );
          }
          else
          {
            leavesMG.Add( leafGeometry );
          }

          

        }

        
      }

      


      if ( ! useDebugSingleMeshGeneration )
      {
        var leavesInstance = this.CreateChild<MeshInstance3D>( "Leaves Mesh" );
        leavesInstance.Mesh = leavesMG.GenerateMesh();
        leavesInstance.SetSurfaceOverrideMaterial( 0, leavesMaterial );
        
        var meshInstance3D = this.CreateChild<MeshInstance3D>( "Tree Mesh");
        meshInstance3D.Mesh = mg.GenerateMesh();
        meshInstance3D.SetSurfaceOverrideMaterial( 0, branchMaterial );
      } 



      if ( hideStructureWhenMeshWasBuild )
      {
        structure.Disable();
      }
  
      numTriangles = mg.numTriangles;

      return;
    }

    [Export]
    public int leavesLevel = 3;

    [Export]
    public Vector3 leavesEulerRotationPre =Vector3.Zero;

    [Export]
    public Vector3 leavesEulerRotationPost =Vector3.Zero;

    [Export]
    public float leafSize = 1;

    [Export]
    public Curve leafSizeScaleVariance;

    [Export]
    public Curve leavesDensity;

    [Export]
    public float leavesNormalBlending = 0.5f;

    [Export]
    public Curve leavesBlendCurve;
    
    [Export]
    public float leavesStartBlend = 4;

    [Export]
    public float leavesEndBlend = 8;

    [Export]
    public int numTriangles = 0;

    [Export]
    public float processTimePerFrame = 1f/80f;

    [Export]
    public float smallestRadius = 0.02f;

    float SmallestChildRadius( TreeBranch tb )
    {
      var radius = smallestRadius;

      if ( tb.children.Length > 0 )
      {
        radius = _branches[ tb.children[ 0 ] ].radius;
      }

      for ( int i = 1; i < tb.children.Length; i++ )
      {
        var childIndex = tb.children[ i ];
        var childTB = _branches[ childIndex ];
        radius = Mathf.Min( radius, childTB.radius );        
      }

      return radius;
    }

    MeshGeometry CreateBranchMesh( TreeBranch tb )
    {
      var tubeGeometry = new TubeGeometry();
      tubeGeometry.curve = tb.GetSpline( this ).GetCurve(); 
      var normalizedLevel = tb.level / (float)( maxLevel );
      var levelSegments = radialSegementsMultiplier.Sample( normalizedLevel );

      branchSettings.radialSegments = Mathf.RoundToInt( radialSegments * levelSegments );
      tubeGeometry.settings = branchSettings;
      tubeGeometry.shapes = shapes;
      branchSettings.radius = 1.0f;

      if ( tb.numRadi == 3 )
      {
        
        tb.curve = MathX.CurveFromPoints( tb.radius, tb.radius2, tb.radius3 );
        branchSettings.radiusSizeCurve = tb.curve;
      }
      else
      {
        branchSettings.radiusSizeCurve = MathX.Curve( tb.radius, SmallestChildRadius( tb ) );
      }

      

      var mesh = tubeGeometry.CreateMesh();
      // mesh.SmoothNormals();  

      

      return mesh;
    }

    [Export]
    public Node3D structure;


    OcTree<TreeBranch> ocTree;
    CustomTreeWalker<TreeBranch> walker;

    async Task GenerateBranches()
    { 
      if ( walker == null )
      {
        walker = new CustomTreeWalker<TreeBranch>(
           tb => _branches[ tb.parent ],
           ( tb, childIndex ) => _branches[ childIndex ],
           tb => tb.children.Count() 
        );
      }
      
      if ( randomizeSeed )
      {
        seed = GodotRandom.Get().IntegerInclusive( 0, 10000 );
      }

      branches = [];
      _branches.Clear();
      this.DestroyChildren();

      
      ocTree = new OcTree<TreeBranch>( GlobalPosition, rootHeight * 3 * maxLevel, 10 );

      structure = this.CreateChild<Node3D>( "Tree Structure" );

      random = LCG.WithSeed( seed );

      branches = _branches.ToArray();

      var nextParents = new List<TreeBranch>();

      var root = GenerateBranch( -1, 0, 0 );
      nextParents.Add( root );

      var time = Async.StartTimer();

      maxChecksFull = 0;

      for ( int i = 0; i < maxLevel; i++ )
      {
        
        var parents = nextParents;
        nextParents = new List<TreeBranch>();

        var normalizedLevel = ( ( (float) i ) / (float) ( maxLevel - 1f ) );

        // this.LogInfo( "Level", i, normalizedLevel, "Parents", parents.Count );
        
        for ( int j =0 ; j < parents.Count; j++ )
        {
          var p = parents[ j ];
          random.SetSeed( p.startSeed + 1234 );
          var numKids = Mathf.RoundToInt( levelToChildBranches.Sample( normalizedLevel ) );

          // this.LogInfo( i, "Parent", j, "Index:", p.index, "Kids:", numKids );

          for ( int k = 0; k < numKids; k++ )
          {
            time = await Async.WaitIfExceeded( time, this, processTimePerFrame );
            nextParents.Add( GenerateBranch( p.index, normalizedLevel, numKids ) );
          }

          // this.LogInfo( "Num Next", nextParents.Count );
        }

      }

      for ( int i = 0; i < _branches.Count; i++ )
      {
        time = await Async.WaitIfExceeded( time, processTimePerFrame );

        var p = _branches[ i ];
        
        if ( p.children.Length == 0 )
        {
          continue;
        }

        var child = walker.GetDirectChildWithHighestValue( p, tb => tb.branchPosition );

        if ( child == null || child.parent != p.index )
        {
          continue;
        }

        var spline = p.GetSpline( this );
        var lastPoint = spline.GetLastChild<SplinePoint>();

        if ( lastPoint == null )
        {
          this.LogInfo( "has no lastPoint:", p.index );
          continue;
        }

        var childSpline = child.GetSpline( this );       
        var firstPoint = childSpline.GetChild<SplinePoint>();
        var lastPoint2 = childSpline.GetLastChild<SplinePoint>();

        p.radius2 = SmallestChildRadius( p );
        p.radius3 = SmallestChildRadius( child ) * 0.5f;
        p.numRadi = 3;


        lastPoint.GlobalPosition = firstPoint.GlobalPosition;
        lastPoint.SetGlobalQuaternion( firstPoint.GetGlobalQuaternion() );

        var additional = spline.CreateChild<SplinePoint>();
        additional.GlobalPosition = lastPoint2.GlobalPosition;
        additional.SetGlobalQuaternion( lastPoint2.GetGlobalQuaternion() );

        spline.ClearCurveCache();
        spline.AutoOrientate();
        spline.GetCurve();
      }

      if ( levelToSubdivisionNoise != null )
      {
        for ( int bb = 0; bb < _branches.Count; bb++ )
        {
          time = await Async.WaitIfExceeded( time, processTimePerFrame );

          var b = _branches[ bb ];

          var normalizedLevel = b.level / maxLevel;
          var subdivisions = Mathf.Ceil( levelToSubdivisionNoise.Sample( normalizedLevel ) );
          var spline = b.GetSpline( this );
          var curve = spline.GetCurve();
          var numPoints = curve.points.Count + subdivisions;

          spline.DestroyChildren();
          spline.ClearCurveCache();
          
          if ( useDebugSingleMeshGeneration && bb % 5 == 0 )
          {
            await this.RequestNextFrame();
          }

          for ( int i = 0; i < numPoints; i++ )
          {              
            time = await Async.WaitIfExceeded( time, processTimePerFrame );

            var t = i / (numPoints - 1f  );
            var t2 = 1f - t;
            var minT = Mathf.Min( t * 2f, t2 * 2f );
            var p = curve.PoseAt( t );
            var rr = random.Sample( subdivisionNoiseRange );
            p.position += random.InSphere( rr * minT * b.height / (float)subdivisions );
            var sp = spline.CreateChild<SplinePoint>();
            sp.editorSplinePointSize = 0.001f;

            p.Set( sp );           
            
          }
          

          var ss = b.GetSpline( this );
          ss.AutoOrientate();
          

        }
      }


      branches = _branches.ToArray();

      if ( hideStructureWhenMeshWasBuild )
      {
        structure.Disable();
      }
  

      return;
    }

    
   

    [Export]
    int maxChecksFull = 0;
    
    TreeBranch GenerateBranch( int parent, float normalizedLevel, int numSiblings )
    {
      
      TreeBranch parentBranch = null;
      var childIndex = 0;

      if ( parent != -1 )
      {
        parentBranch = _branches[ parent ];
        childIndex = parentBranch.children.Length;
      }

      var treeBranch = new TreeBranch();
      treeBranch.parent = parent;
      treeBranch.level  = parentBranch == null ? 0 : parentBranch.level + 1;
      treeBranch.index = _branches.Count;
      treeBranch.startSeed = seed * 100000 + treeBranch.level * 1000 + childIndex;



      random.SetSeed( treeBranch.startSeed );

      if ( parentBranch == null )
      {
        treeBranch.radius = rootRadius;
        treeBranch.height = rootHeight;
      }
      else
      {
        treeBranch.radius = parentBranch.radius * random.Sample( radiusScale );
        treeBranch.height = parentBranch.height * random.Sample( heightScale );  

        parentBranch.children = Arrays.Add( parentBranch.children, treeBranch.index );
      }

      treeBranch.branchPosition = random.SelectCurveWeights( branchPosition, 100 );

      var spline = structure.CreateChild<Spline>( "Branch - " + treeBranch.level + " - " + childIndex);
  

      treeBranch.spline = spline.GetNodePath();

      if ( parentBranch == null )
      {
        var root = spline.CreateChild<SplinePoint>();
        var end  = spline.CreateChild<SplinePoint>();

        root.editorSplinePointSize = 0.001f;

        end.Position = new Vector3( 0, rootHeight, 0 ) + random.InSphere( noise );
        

        spline.autoOrientationUp = Vector3.Back;
      }
      else
      {
        var parentSpline = parentBranch.GetSpline( this );
        var parentCurve = parentSpline.GetCurve();

        
        var rootPose = parentCurve.PoseAt( treeBranch.branchPosition );

        var siblingRotation = 360f * ( childIndex / (float)numSiblings );
        siblingRotation += random.Sample( siblingRotationVariance, 0, 0 );
        
        rootPose.ApplyTwist( siblingRotation );
        var rootRotationCenter = rotationPerLevel.Sample( normalizedLevel );
        var rootRotationRange  = rotationVariation == null ? 0 : random.Sample( rotationVariation );
        rootPose.RotateAround( Math3D.RotateXDegrees( random.Range( rootRotationCenter - rootRotationRange, rootRotationCenter + rootRotationRange ) ), rootPose.position );
        

        var endPose  = rootPose.Clone(); 
        endPose.position = rootPose.position + rootPose.forward * treeBranch.height + random.InSphere( noise );;

        var rootPoint = spline.CreateChild<SplinePoint>();                
        var endPoint = spline.CreateChild<SplinePoint>();

        rootPose.Set( rootPoint );
        endPose.Set( endPoint );

        var checkForOverlaps = true;
        var maxChecks = 30;
        var checks = 0;

        while ( avoidClusters && checkForOverlaps && checks < maxChecks )
        {
          var line = new Line3( rootPoint.GlobalPosition, endPoint.GlobalPosition );

          var testBox = line.GetBox( treeBranch.radius );

          var overlapping = ocTree.GetInsideBox( testBox );

          overlapping = overlapping.Filter( o => o != parentBranch );

          var conflictElement = overlapping.Find(
            ( tb )=>
            {
              var tbLine = tb.GetLine3();

              if ( line.DistanceToLine( tbLine ) < treeBranch.radius * minClusterDistance ) 
              {
                return true;
              }

              return false;
            }
          );

          if ( conflictElement == null )
          {
            checkForOverlaps = false;
          }
          else
          {
            treeBranch.branchPosition = random.SelectCurveWeights( branchPosition, 100 );

            rootPose = parentCurve.PoseAt( treeBranch.branchPosition );
            
            rootPose.ApplyTwist( random.Range( 0, 360 ) );
            rootRotationCenter = rotationPerLevel.Sample( normalizedLevel );
            rootRotationRange  = rotationVariation == null ? 0 : random.Sample( rotationVariation );
            rootPose.RotateAround( Math3D.RotateXDegrees( random.Range( rootRotationCenter - rootRotationRange, rootRotationCenter + rootRotationRange ) ), rootPose.position );
            
            endPose  = rootPose.Clone(); 
            endPose.position = rootPose.position + rootPose.forward * treeBranch.height + random.InSphere( noise );

            rootPose.Set( rootPoint );
            endPose.Set( endPoint );
          }

          checks ++;
          
        }

        if ( checks == maxChecks )
        {
          maxChecksFull ++;
        }

       
      }   

      spline.AutoOrientate();
      spline.SetEditorPointSize( 0.1f / 3f );

      var curve = spline.GetCurve();
      treeBranch.start = curve.PositionAt( 0 );
      treeBranch.end  = curve.PositionAt( 1 );


      var box = spline.GetBounds();

      ocTree.BoxInsert( treeBranch, box );

      _branches.Add( treeBranch );

      return treeBranch;
    }


  }
}