# Copyright (c) Microsoft Corporation. 
# Licensed under the MIT license. 

import math
import os
import sys
import glob
import gc
import threading

from pathlib import Path
from simplygon import simplygon_loader
from simplygon import Simplygon


def AmbientOcclusionCasting(sg: Simplygon.ISimplygon):
    sgSceneImporter = sg.CreateSceneImporter()
    sgSceneImporter.SetImportFilePath( "../Assets/RiggedSimplygonMan/RiggedSimplygonMan.fbx" )
    if not sgSceneImporter.RunImport():
        raise Exception("Failed to load RiggedSimplygonMan/RiggedSimplygonMan.fbx.")
    sgScene = sgSceneImporter.GetScene()

    # Create the aggregation processor. 
    sgAggregationProcessor = sg.CreateAggregationProcessor()

    sgAggregationProcessor.SetScene( sgScene )

    sgAggregationSettings = sgAggregationProcessor.GetAggregationSettings()
    sgMappingImageSettings = sgAggregationProcessor.GetMappingImageSettings()

    # Merge all geometries into a single geometry. 
    sgAggregationSettings.SetMergeGeometries( True )

    # Generates a mapping image which is used after the aggregation to cast new materials to the new 
    # aggregated object. 
    sgMappingImageSettings.SetGenerateMappingImage( True )
    sgMappingImageSettings.SetApplyNewMaterialIds( True )
    sgMappingImageSettings.SetGenerateTangents( True )
    sgMappingImageSettings.SetUseFullRetexturing( True )

    sgOutputMaterialSettings = sgMappingImageSettings.GetOutputMaterialSettings(0)

    # Setting the size of the output material for the mapping image. This will be the output size of the 
    # textures when we do material casting in a later stage. 
    sgOutputMaterialSettings.SetTextureWidth( 2048 )
    sgOutputMaterialSettings.SetTextureHeight( 2048 )

    # Start the aggregation process. 
    sgAggregationProcessor.RunProcessing()

    # Setup and run the ambient occlusion material casting. 
    sgAmbientCaster = sg.CreateAmbientOcclusionCaster()
    sgAmbientCaster.SetMappingImage( sgAggregationProcessor.GetMappingImage() )
    sgAmbientCaster.SetSourceMaterials( sgScene.GetMaterialTable() )
    sgAmbientCaster.SetSourceTextures( sgScene.GetTextureTable() )
    sgAmbientCaster.SetOutputFilePath( "AmbientTexture" )

    sgAmbientCasterSettings = sgAmbientCaster.GetAmbientOcclusionCasterSettings()
    sgAmbientCasterSettings.SetMaterialChannel( "Ambient" )
    sgAmbientCasterSettings.SetRaysPerPixel( 64 )
    sgAmbientCasterSettings.SetOutputImageFileFormat( Simplygon.EImageOutputFormat_PNG )

    sgAmbientCaster.RunProcessing()
    ambientTextureFilePath = sgAmbientCaster.GetOutputFilePath()

    # Update scene with new casted texture. 
    sgMaterialTable = sg.CreateMaterialTable()
    sgTextureTable = sg.CreateTextureTable()
    sgMaterial = sg.CreateMaterial()
    sgAmbientTexture = sg.CreateTexture()
    sgAmbientTexture.SetName( "Ambient" )
    sgAmbientTexture.SetFilePath( ambientTextureFilePath )
    sgTextureTable.AddTexture( sgAmbientTexture )

    sgAmbientTextureShadingNode = sg.CreateShadingTextureNode()
    sgAmbientTextureShadingNode.SetTexCoordLevel( 0 )
    sgAmbientTextureShadingNode.SetTextureName( "Ambient" )

    sgMaterial.AddMaterialChannel( "Ambient" )
    sgMaterial.SetShadingNetwork( "Ambient", sgAmbientTextureShadingNode )

    sgMaterialTable.AddMaterial( sgMaterial )

    sgScene.GetTextureTable().Clear()
    sgScene.GetMaterialTable().Clear()
    sgScene.GetTextureTable().Copy(sgTextureTable)
    sgScene.GetMaterialTable().Copy(sgMaterialTable)

    sgSceneExporter = sg.CreateSceneExporter()
    sgSceneExporter.SetScene(sgScene)
    sgSceneExporter.SetExportFilePath( "Output.fbx" )
    if not sgSceneExporter.RunExport():
        raise Exception("Failed to save Output.fbx.")


if __name__ == '__main__':
    sg = simplygon_loader.init_simplygon()

    AmbientOcclusionCasting(sg)

    sg = None
    gc.collect()

