rj-action-library/Runtime/Procedural/Baking/MultiBaker/MultiBaker.cs

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

using System.Threading.Tasks;
using System.Linq;

namespace Rokojori
{
  [Tool]
  [GlobalClass]
  public partial class MultiBaker:Node
  {
    [ExportToolButton( "Bake")]
    public Callable BakeButton => Callable.From( () => StartBaking() );

    [Export]
    public bool cleanUpAfterBaking = true;

    [Export]
    public _XX_MultiBakeMode bakeMode;

    public enum MaterialMode
    {
      Full_Seperated, 
      Simple_Prebaked
    }

    // [ExportGroup("Textures")]

    [Export]
    public _XX_MultiTextureBaker textureBaker;


    // [Export]
    // public MaterialMode materialMode;

    // [Export]
    // public int dilationRadius = 64;
    // [ExportGroup("Material/Full Seperated")]
    // [Export]
    // public bool mmfs_Normals = true; 
    // [Export]
    // public bool mmfs_Depth = true;
    // [Export]
    // public bool mmfs_ORM = true;

    [Export]
    public BakingViewSettings viewSettings;

    [ExportGroup( "Object")]
    
    [Export]
    public Node3D target;
    
    [Export]
    public bool autoTargetPivot = false;

    [Export]
    public Vector3 targetPivot;

    

    

    public float cameraZoom
    {
      get 
      {
        if ( viewSettings.fovDistance is AutoDistance_BakingFDSettings ad )
        {
          return ad.zoom;
        }
        
        if ( viewSettings.fovDistance is AutoDistance_BakingFDSettings afd )
        {
          return afd.zoom;
        }


        return 1;
      }
    }

    [ExportGroup( "Output")]


    [Export]
    public string outputDirectory;

    [Export]
    public string outputFileName;

    [Export]
    public float outputQuality = 1f;

    [Export]
    public Vector2 outputTextureSize = new Vector2( 2048, 2048 );    

    [Export]
    public bool showOutputTexture = false;


    [ExportGroup("Read Only")]
    [Export]
    public SubViewport X_bakingViewport;

    [Export]
    public Node3D X_bakingTargetContainer;

    [Export]
    public Node X_views;

    [Export]
    public WorldEnvironment X_worldEnvironment;    
  
    [Export]
    public MeshInstance3D X_outputMesh;

    [Export]
    public TextureMerger X_textureMerger;

    [Export]
    public TextureDilate X_textureDilate;

    [Export]
    public SetBakingMaterials X_setBakingMaterials;

    [Export]
    public MeshInstance3D X_texturePreview;

    [Export]
    public CsgMesh3D X_depthMapper;

    [Export]
    public Texture2D X_bakedTextureAlbedo;

    [Export]
    public Texture2D X_bakedTextureNormal;

    [Export]
    public Texture2D X_bakedTextureORM;

    [Export]
    public Texture2D X_bakedTextureDepth;



    [Export]
    public bool X_baking = false;

    public void StartBaking()
    {
      Initialize();
      Bake();
    }

    public void Initialize()
    {     
      this.LogInfo( "Initializing" );
      
      Nodes.RemoveAndDeleteChildren( this );

      _bakerCameras = null;

      X_bakingViewport = this.CreateChild<SubViewport>( "Multi Baker Viewport" );

      X_bakingViewport.Size = (Vector2I) outputTextureSize;
      X_bakingViewport.OwnWorld3D = true;
      X_bakingViewport.TransparentBg = true;

      X_worldEnvironment = X_bakingViewport.CreateChild<WorldEnvironment>( "Multi Baker Environment" );
      X_worldEnvironment.Environment = new Godot.Environment();
      X_worldEnvironment.Environment.AmbientLightSource = Godot.Environment.AmbientSource.Color;
      X_worldEnvironment.Environment.AmbientLightColor = HSLColor.white;

      X_bakingTargetContainer = X_bakingViewport.CreateChild<Node3D>( "Target Container" );

      X_views = X_bakingViewport.CreateChild<Node>( "Views" );

      X_textureMerger = this.CreateChild<TextureMerger>( "Texture Merger" );
      X_textureMerger.multiBaker = this;
      // X_textureMerger.dilationRadius = dilationRadius;
      X_textureMerger.Initialize();

      X_textureDilate = this.CreateChild<TextureDilate>( "Texture Dilate" );
      X_textureDilate.viewport = X_textureMerger.X_textureMergerViewport;

      X_outputMesh = this.CreateChild<MeshInstance3D>( "Output Mesh" );

      X_texturePreview = this.CreateChild<MeshInstance3D>( "Texture Preview" );
      X_texturePreview.Mesh = new QuadMesh();
      X_texturePreview.Scale = Vector3.One * 100;

      var pm = new StandardMaterial3D();
      pm.Transparency = BaseMaterial3D.TransparencyEnum.AlphaScissor;
      pm.ResourceLocalToScene = true;

      var vt = new ViewportTexture();
      vt.ViewportPath = X_textureMerger.X_textureMergerViewport.GetPath();
      pm.AlbedoTexture = vt;

      X_setBakingMaterials = this.CreateChild<SetBakingMaterials>( "Set Baking Materials" );

      Materials.Set( X_texturePreview, pm );                  

    }

    public async Task<Texture2D> GrabDilatedTexture( bool srgb, bool alpha, bool mipmaps )
    {
      X_textureDilate.viewport = X_textureMerger.X_textureMergerViewport;
      var maxRadius = X_textureMerger.GetMaxTextureDimension() / 2;
      X_textureDilate.SetDilationRadius( maxRadius );
      var texture = await X_textureDilate.Grab( srgb, alpha, mipmaps );

      return texture;
    }   

    bool _targetIsInInitialState = false;

    public void SetTargetDirty()
    {
      _targetIsInInitialState = false;
    }

    public void ResetOriginalTargetState()
    {
      if ( _targetIsInInitialState )
      {
        return;
      }

      // _bakerCameras.ForEach( b => b.viewport.en)

      Nodes.RemoveAndDeleteChildren( X_bakingTargetContainer );
      target.DeepCopyTo( X_bakingTargetContainer );

      _targetIsInInitialState = true;

      return;
    }


    bool _compositorsClean = false;

    public void SetCompositorsDirty()
    {
      _compositorsClean = false;
    }


    public void ResetCompositorStates()
    {
      X_worldEnvironment.Compositor = null;
      _compositorsClean = true;
    }

    public void AddCompositorEffect( CompositorEffect compositorEffect )
    {
      if ( X_worldEnvironment.Compositor == null )
      {
        X_worldEnvironment.Compositor = new Compositor();
      }

      var compositorEffects = X_worldEnvironment.Compositor.CompositorEffects;
      compositorEffects.Add( compositorEffect );

      X_worldEnvironment.Compositor.CompositorEffects = compositorEffects;

      _compositorsClean = false; 
    } 
   

    public async Task Bake()
    {
      if ( X_baking )
      {
        return;
      }

      X_baking = true;


      try
      {
        bakeMode.multiBaker = this;

        _targetIsInInitialState = false;

        this.LogInfo( "Started baking" );

        ResetOriginalTargetState();

        if ( _bakerCameras == null || _bakerCameras.Count != GetNumViews() )
        {
          CreateBakerCameras();
          await this.RequestNextFrame();
        }

        SetupBakerCameras();

        this.LogInfo( "Bake Cameras set up" );
        await this.RequestNextFrame();

        this.LogInfo( "Render passes" );
        var passes = textureBaker.GetPasses( this );


        //  X_setBakingMaterials.SetTarget( X_bakingTargetContainer );

        // var bakingMaterialModes = new List<BakingMaterialMode>();

        // var preview_QuickMaterial = MaterialMode.Simple_Prebaked == materialMode;

        // bakeMode.CreateMaterial( preview_QuickMaterial );

        bakeMode.CreateMaterial( passes );

        // if ( preview_QuickMaterial )
        // { 
        //   bakingMaterialModes.Add( BakingMaterialMode.Preview );
        // }
        // else
        // {
        //   bakingMaterialModes.Add( BakingMaterialMode.Albedo );

        //   if ( mmfs_Normals )
        //   {
        //     bakingMaterialModes.Add( BakingMaterialMode.Normals );
        //   }

        //   if ( mmfs_Depth )
        //   {
        //     bakingMaterialModes.Add( BakingMaterialMode.Depth );
        //   }

        //   if ( mmfs_ORM )
        //   {
        //     bakingMaterialModes.Add( BakingMaterialMode.ORM );
        //   }
        // }

        this.LogInfo( "Prepared baking modes" );
        X_textureMerger.textureSize = outputTextureSize;
        X_textureMerger.Initialize();
        X_textureMerger.CreateLayout();


        // this.LogInfo( "Prepared texture merger" );
        // var fovDistance = viewSettings.fovDistance.ComputeFOVDistance( 
        // var objectDistance = GetCameraDistance();

        // this.LogInfo( "Set Camera Distance", objectDistance );


        for ( int i = 0; i < passes.Count; i++ )
        { 
          ResetOriginalTargetState();
          
          _bakerCameras.ForEach( b => b.ApplyCameraSettings() );
          await passes[ i ].Bake();
          await this.RequestNextFrame();

        }

        bakeMode.AssignMaterial( passes );

        // for ( int i = 0; i < bakingMaterialModes.Count; i++ )
        // {
        //   this.LogInfo( "Baking mode:", bakingMaterialModes[ i ] );   
        //   X_setBakingMaterials.mode = bakingMaterialModes[ i ];
        //   X_setBakingMaterials.ApplyBakingMaterials( objectDistance, _targetBoundingSphere.radius );

        //   this.LogInfo( "Materials changed:", bakingMaterialModes[ i ] );

        //   _bakers.ForEach( b => b.ApplyCameraSettings() );

        //   await this.RequestNextFrame();

        //   X_textureDilate.viewport = X_textureMerger.X_textureMergerViewport;
        //   Texture2D texture = await X_textureDilate.Grab();

          
        //   this.LogInfo( "Assigning Texture", bakingMaterialModes[ i ] );
        //   bakeMode.AssignMaterial( bakingMaterialModes[ i ], texture ); 

        //   this.LogInfo( "Baking done:", bakingMaterialModes[ i ] );

        //   await this.RequestNextFrame();

        // }    

        // await this.RequestNextFrame();

        X_outputMesh.GlobalPosition = target.GlobalPosition - targetPivot;

        X_outputMesh.Scale = Vector3.One * cameraZoom;

       
        this.LogInfo( "All Baking done" );

      }
      catch ( System.Exception e )
      {
        this.LogError( "Baking failed" );
        this.LogError( e );
      }

      if ( cleanUpAfterBaking )
      {
        this.LogInfo( "Cleaning Up" );
        CleanUp();
      }


      X_baking = false;
      return;

    }
    
    public List<SubViewport> GetAllViewports()
    {
      return Nodes.AllIn<SubViewport>( X_views, null, false );
    }

    // public void CacheTexture( BakingMaterialMode mode, Texture2D texture )
    // {
    //   if ( BakingMaterialMode.Albedo == mode )
    //   {
    //     X_bakedTextureAlbedo = texture;
    //   }
    //   else if ( BakingMaterialMode.Normals == mode )
    //   {
    //     X_bakedTextureNormal = texture;
    //   }
    //   else if ( BakingMaterialMode.ORM == mode )
    //   {
    //     X_bakedTextureORM = texture;
    //   }
    //   else if ( BakingMaterialMode.Depth == mode )
    //   {
    //     X_bakedTextureDepth = texture;
    //   }

    // } 

    public void CleanUp()
    {
      var mesh = X_outputMesh;
      mesh.Reparent( GetParent(), true );

      X_bakingViewport = null;

      X_bakingTargetContainer = null;

      X_views = null;

      X_worldEnvironment = null;    
    
      X_outputMesh = null;

      X_textureMerger = null;

      X_textureDilate = null;

      X_setBakingMaterials = null;

      X_texturePreview = null;

      X_depthMapper = null;

      X_bakedTextureAlbedo = null;

      X_bakedTextureNormal = null;

      X_bakedTextureORM = null;

      X_bakedTextureDepth = null;

      Nodes.RemoveAndDeleteChildren( this );

      mesh.Reparent( this, true );
    }

    

    public int GetNumViews()
    {
      return bakeMode.GetNumViews();
    }

    List<Baker> _bakerCameras;

    public List<Baker> bakerCameras => _bakerCameras;

    public void CreateBakerCameras()
    {    

      Nodes.RemoveAndDeleteChildren( X_views );

      var numViews = GetNumViews();

      _bakerCameras = new List<Baker>();

      var alginment = TextureMerger.ComputeTextureAlignment( numViews );
      var minViewsPerAxis = Mathf.Max( alginment.X, alginment.Y );

      this.LogInfo( "Size per Baker:", "nv", numViews, "al", alginment, "mvpa", minViewsPerAxis,"s", (Vector2I) ( outputTextureSize / minViewsPerAxis ) );
      
      for ( int i = 0; i < numViews; i++ )
      {
        var userIndex = ( i + 1 );
        var bakingView = X_views.CreateChild<SubViewport>( "Baking View " + userIndex );
        bakingView.TransparentBg = true;
        bakingView.Size = (Vector2I) ( outputTextureSize / minViewsPerAxis );
        
        var bakingCamera = bakingView.CreateChild<Camera3D>( "Camera View " + userIndex );
        var baker = bakingView.CreateChild<Baker>( "Baker " + userIndex );

        baker.camera   = bakingCamera;
        baker.target   = X_bakingTargetContainer;
        baker.viewport = bakingView;
        baker.viewport.Msaa3D = Viewport.Msaa.Msaa8X;
        


        

        _bakerCameras.Add( baker );
      }
    }

    public void SetupBakerCameras()
    {
      var numViews = GetNumViews();
      var minViewsPerAxis = TextureMerger.ComputeTextureAlignment( numViews ).Y;

      X_bakingViewport.Size = (Vector2I) outputTextureSize;

      for ( int i = 0; i < numViews; i++ )
      {
        var baker = _bakerCameras[ i ];
        var bakingView = baker.viewport as SubViewport;
        bakingView.Size = (Vector2I) ( outputTextureSize / minViewsPerAxis );
      }

      _targetBoundingSphere = null;
      _targetBoundingBox = null;

      ComputeBoundingSphere();

      if ( _targetBoundingSphere == null )
      {
        this.LogError( "No bounding sphere created, ensure there are visible targets" );
        return;
      }

      ComputeCameraViewSettings();

      bakeMode.CreateBakers();

      

    }

    Sphere _targetBoundingSphere;
    Box3 _targetBoundingBox;

    Sphere targetBoundingSphere 
    {
      get 
      {
        if ( _targetBoundingSphere == null )
        {
          ComputeBoundingSphere();
        }

        return _targetBoundingSphere;
      }
    }

    void ComputeBoundingSphere()
    {
      if ( X_bakingTargetContainer.GetChildCount() == 0 )
      {
        _targetBoundingSphere = null;
        return;
      }

      var firstChild = X_bakingTargetContainer.GetChild( 0 ) as Node3D;

      if ( firstChild == null )
      {
        _targetBoundingSphere = null;
        return;
      }

      if ( ! firstChild.Visible )
      {
        firstChild.Visible = true;
      }

      _targetBoundingBox = X_bakingTargetContainer.GetWorldBounds();

      if ( _targetBoundingBox == null )
      {
        return;
      }

      if ( autoTargetPivot )
      { 
        Box3 box = target.GetWorldBounds();

        if ( box == null )
        {
          _targetBoundingSphere = null;
          return;
        }

        targetPivot = new Vector3( 0, -box.center.Y, 0 );
      }
      
      _targetBoundingSphere = Sphere.ContainingBox( _targetBoundingBox );
    }

    float _cameraFOV = 0;
    float _cameraDistance = 0;
    float _outputScale = 1;
    
    void ComputeCameraViewSettings()
    {
      var fovDistance = viewSettings.fovDistance;
      var size = targetBoundingSphere.radius;

      var computeScale = false;

      if ( fovDistance is AutoDistance_BakingFDSettings ad )
      {
        size = ad.sizeEstimation == _XX_BakingFDSettings.SizeEstimationType.Bounding_Sphere ? 
                targetBoundingSphere.radius : ( _targetBoundingBox.size.Y / 2f );
        computeScale = true;
      } 

      if ( fovDistance is AutoFOVDistance_BakingFDSettings afd )
      {
        size = afd.sizeEstimation == _XX_BakingFDSettings.SizeEstimationType.Bounding_Sphere ? 
                targetBoundingSphere.radius : ( _targetBoundingBox.size.Y / 2f );
        computeScale = true;
      } 
      
      var fd = fovDistance.ComputeFOVDistance( size / 2f );

      _cameraFOV = fd.X;
      _cameraDistance = fd.Y;
      _outputScale = computeScale ? Cameras.ComputeCameraFittingScale( _cameraFOV, _cameraDistance ) : 1;
    }
    
    public float GetCameraFOV()
    {
      return _cameraFOV;
    }

    public float GetCameraDistance()
    {
      return _cameraDistance;
    }

    public float GetOutputScale()
    {
      return _outputScale;
    }
  }
}