o3de/Gems/SceneProcessing/Code/Source/Generation/Components/TangentGenerator/TangentGenerateComponent.cpp

/*
 * Copyright (c) Contributors to the Open 3D Engine Project.
 * For complete copyright and license terms please see the LICENSE at the root of this distribution.
 *
 * SPDX-License-Identifier: Apache-2.0 OR MIT
 *
 */

#include <Generation/Components/TangentGenerator/TangentGenerateComponent.h>
#include <Generation/Components/TangentGenerator/TangentGenerators/MikkTGenerator.h>
#include <Generation/Components/TangentGenerator/TangentGenerators/BlendShapeMikkTGenerator.h>

#include <SceneAPI/SceneCore/DataTypes/Groups/IGroup.h>
#include <SceneAPI/SceneCore/DataTypes/GraphData/IMeshVertexUVData.h>
#include <SceneAPI/SceneCore/DataTypes/GraphData/IMeshVertexTangentData.h>
#include <SceneAPI/SceneCore/DataTypes/GraphData/IMeshVertexBitangentData.h>

#include <SceneAPI/SceneData/Rules/TangentsRule.h>

#include <SceneAPI/SceneCore/Containers/Views/PairIterator.h>
#include <SceneAPI/SceneCore/Containers/Views/SceneGraphDownwardsIterator.h>
#include <SceneAPI/SceneCore/Containers/Views/SceneGraphChildIterator.h>
#include <SceneAPI/SceneCore/Containers/Views/ConvertIterator.h>
#include <SceneAPI/SceneCore/Containers/Utilities/Filters.h>
#include <SceneAPI/SceneCore/Utilities/Reporting.h>
#include <SceneAPI/SceneCore/DataTypes/DataTypeUtilities.h>
#include <SceneAPI/SceneCore/DataTypes/GraphData/IMeshData.h>
#include <SceneAPI/SceneData/GraphData/BlendShapeData.h>
#include <SceneAPI/SceneData/GraphData/MeshVertexBitangentData.h>
#include <SceneAPI/SceneData/GraphData/MeshVertexTangentData.h>

#include <AzToolsFramework/Debug/TraceContext.h>

#include <AzCore/Math/Vector4.h>
#include <AzCore/std/smart_ptr/make_shared.h>


namespace AZ::SceneGenerationComponents
{
    TangentGenerateComponent::TangentGenerateComponent()
    {
        BindToCall(&TangentGenerateComponent::GenerateTangentData);
    }


    void TangentGenerateComponent::Reflect(AZ::ReflectContext* context)
    {
        AZ::SerializeContext* serializeContext = azrtti_cast<AZ::SerializeContext*>(context);
        if (serializeContext)
        {
            serializeContext->Class<TangentGenerateComponent, AZ::SceneAPI::SceneCore::GenerationComponent>()->Version(1);
        }
    }

    AZ::SceneAPI::DataTypes::TangentSpace TangentGenerateComponent::GetTangentSpaceFromRule(const AZ::SceneAPI::Containers::Scene& scene) const
    {
        for (const auto& object : scene.GetManifest().GetValueStorage())
        {
            if (object->RTTI_IsTypeOf(AZ::SceneAPI::DataTypes::IGroup::TYPEINFO_Uuid()))
            {
                const AZ::SceneAPI::DataTypes::IGroup* group = azrtti_cast<const AZ::SceneAPI::DataTypes::IGroup*>(object.get());
                const AZ::SceneAPI::SceneData::TangentsRule* rule = group->GetRuleContainerConst().FindFirstByType<AZ::SceneAPI::SceneData::TangentsRule>().get();
                if (rule)
                {
                    return rule->GetTangentSpace();
                }
            }
        }

        return AZ::SceneAPI::DataTypes::TangentSpace::FromSourceScene;
    }

    AZ::SceneAPI::Events::ProcessingResult TangentGenerateComponent::GenerateTangentData(TangentGenerateContext& context)
    {
        // Iterate over all graph content and filter out all meshes.
        AZ::SceneAPI::Containers::SceneGraph& graph = context.GetScene().GetGraph();
        AZ::SceneAPI::Containers::SceneGraph::ContentStorageData graphContent = graph.GetContentStorage();

        // Build a list of mesh data nodes.
        AZStd::vector<AZStd::pair<AZ::SceneAPI::DataTypes::IMeshData*, AZ::SceneAPI::Containers::SceneGraph::NodeIndex> > meshes;
        for (auto item = graphContent.begin(); item != graphContent.end(); ++item)
        {
            // Skip anything that isn't a mesh.
            if (!(*item) || !(*item)->RTTI_IsTypeOf(AZ::SceneAPI::DataTypes::IMeshData::TYPEINFO_Uuid()))
            {
                continue;
            }

            // Get the mesh data and node index and store them in the vector as a pair, so we can iterate over them later.
            auto* mesh = static_cast<AZ::SceneAPI::DataTypes::IMeshData*>(item->get());
            AZ::SceneAPI::Containers::SceneGraph::NodeIndex nodeIndex = graph.ConvertToNodeIndex(item);
            meshes.emplace_back(mesh, nodeIndex);
        }

        // Iterate over them. We had to build the array before as this method can insert new nodes, so using the iterator directly would fail.
        for (auto& [mesh, nodeIndex] : meshes)
        {
            // Generate tangents for the mesh (if this is desired or needed).
            if (!GenerateTangentsForMesh(context.GetScene(), nodeIndex, mesh))
            {
                return AZ::SceneAPI::Events::ProcessingResult::Failure;
            }

            // Now that we have the tangents and bitangents, calculate the tangent w values for the ones that we imported from the scene file, as they only have xyz.
            UpdateFbxTangentWValues(graph, nodeIndex, mesh);
        }

        return AZ::SceneAPI::Events::ProcessingResult::Success;
    }

    void TangentGenerateComponent::UpdateFbxTangentWValues(AZ::SceneAPI::Containers::SceneGraph& graph, const AZ::SceneAPI::Containers::SceneGraph::NodeIndex& nodeIndex, const AZ::SceneAPI::DataTypes::IMeshData* meshData)
    {
        // Iterate over all UV sets.
        AZ::SceneAPI::DataTypes::IMeshVertexUVData* uvData = FindUvData(graph, nodeIndex, 0);
        size_t uvSetIndex = 0;
        while (uvData)
        {
            AZ::SceneAPI::DataTypes::IMeshVertexTangentData* fbxTangentData = FindTangentData(graph, nodeIndex, uvSetIndex);
            AZ::SceneAPI::DataTypes::IMeshVertexBitangentData* fbxBitangentData = FindBitangentData(graph, nodeIndex, uvSetIndex);

            if (fbxTangentData && fbxBitangentData)
            {
                const size_t numVerts = uvData->GetCount();
                AZ_Assert((numVerts == fbxTangentData->GetCount()) && (numVerts == fbxBitangentData->GetCount()), "Number of vertices inside UV set is not the same as number of tangents and bitangents.");
                for (size_t i = 0; i < numVerts; ++i)
                {
                    // This code calculates the best tangent.w value, which is either -1 or +1, depending on the bitangent being mirrored or not.
                    // We determine this by checking the angle between the generated tangent by doing a cross product between the tangent and normal, and the actual real bitangent.
                    // It is no guarantee that using "cross(normal, tangent.xyz)* tangent.w" will result in the right bitangent, as the basis might not be orthogonal.
                    // But we still go for the best guess.
                    AZ::Vector4 tangent = fbxTangentData->GetTangent(i);
                    AZ::Vector3 tangentDir = tangent.GetAsVector3();
                    tangentDir.NormalizeSafe();
                    AZ::Vector3 normal = meshData->GetNormal(static_cast<AZ::u32>(i));
                    normal.NormalizeSafe();
                    AZ::Vector3 generatedBitangent = normal.Cross(tangentDir);

                    float dot = fbxBitangentData->GetBitangent(i).Dot(generatedBitangent);
                    dot = AZ::GetMax(dot, -1.0f);
                    dot = AZ::GetMin(dot, 1.0f);
                    const float angle = acosf(dot);
                    if (angle > AZ::Constants::HalfPi)
                    {
                        tangent = fbxTangentData->GetTangent(i);
                        tangent.SetW(-1.0f);
                    }
                    else
                    {
                        tangent = fbxTangentData->GetTangent(i);
                        tangent.SetW(1.0f);
                    }
                    fbxTangentData->SetTangent(i, tangent);
                }
            }

            // Find the next UV set.
            uvData = FindUvData(graph, nodeIndex, ++uvSetIndex);
        }
    }

    void TangentGenerateComponent::FindBlendShapes(
        AZ::SceneAPI::Containers::SceneGraph& graph, const AZ::SceneAPI::Containers::SceneGraph::NodeIndex& nodeIndex,
        AZStd::vector<AZ::SceneData::GraphData::BlendShapeData*>& outBlendShapes) const
    {
        const auto nameContentView = AZ::SceneAPI::Containers::Views::MakePairView(graph.GetNameStorage(), graph.GetContentStorage());

        auto meshChildView = AZ::SceneAPI::Containers::Views::MakeSceneGraphChildView<AZ::SceneAPI::Containers::Views::AcceptEndPointsOnly>(
            graph, nodeIndex, nameContentView.begin(), true);
        for (auto child = meshChildView.begin(); child != meshChildView.end(); ++child)
        {
            AZ::SceneData::GraphData::BlendShapeData* blendShape =
                azrtti_cast<AZ::SceneData::GraphData::BlendShapeData*>(child->second.get());
            if (blendShape)
            {
                outBlendShapes.emplace_back(blendShape);
            }
        }
    }

    bool TangentGenerateComponent::GenerateTangentsForMesh(AZ::SceneAPI::Containers::Scene& scene, const AZ::SceneAPI::Containers::SceneGraph::NodeIndex& nodeIndex, AZ::SceneAPI::DataTypes::IMeshData* meshData)
    {
        AZ::SceneAPI::Containers::SceneGraph& graph = scene.GetGraph();

        // Check if we have any UV data, if not, we cannot possibly generate the tangents.
        const size_t uvSetCount = CalcUvSetCount(graph, nodeIndex);
        if (uvSetCount == 0)
        {
            AZ_Warning(AZ::SceneAPI::Utilities::WarningWindow, false, "Cannot generate tangents for this mesh, as it has no UV coordinates.\n");
            return true; // No fatal error
        }

        // Check what tangent spaces we need.
        const AZ::SceneAPI::DataTypes::TangentSpace ruleTangentSpace = GetTangentSpaceFromRule(scene);

        // Find all blend shape data under the mesh. We need to generate the tangent and bitangent for blend shape as well.
        AZStd::vector<AZ::SceneData::GraphData::BlendShapeData*> blendShapes;
        FindBlendShapes(graph, nodeIndex, blendShapes);

        // Generate tangents/bitangents for all uv sets.
        bool allSuccess = true;
        for (size_t uvSetIndex = 0; uvSetIndex < uvSetCount; ++uvSetIndex)
        {
            AZ::SceneAPI::DataTypes::IMeshVertexUVData* uvData = FindUvData(graph, nodeIndex, uvSetIndex);
            if (!uvData)
            {
                AZ_Error(AZ::SceneAPI::Utilities::ErrorWindow, false, "Cannot generate tangents for uv set %zu as it cannot be retrieved.\n", uvSetIndex);
                continue;
            }

            // Check if we had tangents inside the source scene file.
            AZ::SceneAPI::DataTypes::TangentSpace tangentSpace = ruleTangentSpace;
            AZ::SceneAPI::DataTypes::IMeshVertexTangentData* tangentData = FindTangentData(graph, nodeIndex, uvSetIndex);
            AZ::SceneAPI::DataTypes::IMeshVertexBitangentData* bitangentData = FindBitangentData(graph, nodeIndex, uvSetIndex);

            // If all we need is import from the source scene, and we have tangent data from the source scene already, then skip generating.
            if ((tangentSpace == AZ::SceneAPI::DataTypes::TangentSpace::FromSourceScene))
            {
                if (tangentData && bitangentData)
                {
                    AZ_TracePrintf(AZ::SceneAPI::Utilities::LogWindow, "Using source scene tangents and bitangents for uv set %zu for mesh '%s'.\n",
                        uvSetIndex, scene.GetGraph().GetNodeName(nodeIndex).GetName());
                    continue;
                }
                else
                {
                    // In case there are no tangents/bitangents while the user selected to use the source ones, default to MikkT.
                    AZ_Warning(AZ::SceneAPI::Utilities::WarningWindow, false, "Cannot use source scene tangents as there are none in the asset for mesh '%s' for uv set %zu. Defaulting to generating tangents using MikkT.\n",
                        scene.GetGraph().GetNodeName(nodeIndex).GetName(), uvSetIndex);
                    tangentSpace = AZ::SceneAPI::DataTypes::TangentSpace::MikkT;
                }
            }

            if (!tangentData)
            {
                if (!AZ::SceneGenerationComponents::TangentGenerateComponent::CreateTangentLayer(scene.GetManifest(), nodeIndex, meshData->GetVertexCount(), uvSetIndex,
                    tangentSpace, graph, &tangentData))
                {
                    AZ_Error(AZ::SceneAPI::Utilities::ErrorWindow, false, "Failed to create tangents data set for mesh %s for uv set %zu.\n",
                        scene.GetGraph().GetNodeName(nodeIndex).GetName(), uvSetIndex);
                    continue;
                }
            }
            if (!bitangentData)
            {
                if (!AZ::SceneGenerationComponents::TangentGenerateComponent::CreateBitangentLayer(scene.GetManifest(), nodeIndex, meshData->GetVertexCount(), uvSetIndex,
                    tangentSpace, graph, &bitangentData))
                {
                    AZ_Error(AZ::SceneAPI::Utilities::ErrorWindow, false, "Failed to create bitangents data set for mesh %s for uv set %zu.\n",
                        scene.GetGraph().GetNodeName(nodeIndex).GetName(), uvSetIndex);
                    continue;
                }
            }
            tangentData->SetTangentSpace(tangentSpace);
            bitangentData->SetTangentSpace(tangentSpace);

            switch (tangentSpace)
            {
            // Generate using MikkT space.
            case AZ::SceneAPI::DataTypes::TangentSpace::MikkT:
            {
                allSuccess &= AZ::TangentGeneration::Mesh::MikkT::GenerateTangents(meshData, uvData, tangentData, bitangentData);

                for (AZ::SceneData::GraphData::BlendShapeData* blendShape : blendShapes)
                {
                    allSuccess &= AZ::TangentGeneration::BlendShape::MikkT::GenerateTangents(blendShape, uvSetIndex);
                }
            }
            break;

            default:
            {
                AZ_Assert(false, "Unknown tangent space selected (spaceID=%d) for UV set %d, cannot generate tangents!\n", static_cast<AZ::u32>(tangentSpace), uvSetIndex);
                allSuccess = false;
            }
            }
        }

        return allSuccess;
    }

    size_t TangentGenerateComponent::CalcUvSetCount(AZ::SceneAPI::Containers::SceneGraph& graph, const AZ::SceneAPI::Containers::SceneGraph::NodeIndex& nodeIndex) const
    {
        const auto nameContentView = AZ::SceneAPI::Containers::Views::MakePairView(graph.GetNameStorage(), graph.GetContentStorage());

        size_t result = 0;
        auto meshChildView = AZ::SceneAPI::Containers::Views::MakeSceneGraphChildView<AZ::SceneAPI::Containers::Views::AcceptEndPointsOnly>(graph, nodeIndex, nameContentView.begin(), true);
        for (auto child = meshChildView.begin(); child != meshChildView.end(); ++child)
        {
            AZ::SceneAPI::DataTypes::IMeshVertexUVData* data = azrtti_cast<AZ::SceneAPI::DataTypes::IMeshVertexUVData*>(child->second.get());
            if (data)
            {
                result++;
            }
        }

        return result;
    }

    AZ::SceneAPI::DataTypes::IMeshVertexUVData* TangentGenerateComponent::FindUvData(AZ::SceneAPI::Containers::SceneGraph& graph, const AZ::SceneAPI::Containers::SceneGraph::NodeIndex& nodeIndex, AZ::u64 uvSet) const
    {
        const auto nameContentView = AZ::SceneAPI::Containers::Views::MakePairView(graph.GetNameStorage(), graph.GetContentStorage());

        AZ::u64 uvSetIndex = 0;
        auto meshChildView = AZ::SceneAPI::Containers::Views::MakeSceneGraphChildView<AZ::SceneAPI::Containers::Views::AcceptEndPointsOnly>(graph, nodeIndex, nameContentView.begin(), true);
        for (auto child = meshChildView.begin(); child != meshChildView.end(); ++child)
        {
            AZ::SceneAPI::DataTypes::IMeshVertexUVData* data = azrtti_cast<AZ::SceneAPI::DataTypes::IMeshVertexUVData*>(child->second.get());
            if (data)
            {
                if (uvSetIndex == uvSet)
                {
                    return data;
                }
                uvSetIndex++;
            }
        }

        return nullptr;
    }

    AZ::SceneAPI::DataTypes::IMeshVertexTangentData* TangentGenerateComponent::FindTangentData(AZ::SceneAPI::Containers::SceneGraph& graph, const AZ::SceneAPI::Containers::SceneGraph::NodeIndex& nodeIndex, AZ::u64 setIndex) const
    {
        const auto nameContentView = AZ::SceneAPI::Containers::Views::MakePairView(graph.GetNameStorage(), graph.GetContentStorage());

        auto meshChildView = AZ::SceneAPI::Containers::Views::MakeSceneGraphChildView<AZ::SceneAPI::Containers::Views::AcceptEndPointsOnly>(graph, nodeIndex, nameContentView.begin(), true);
        for (auto child = meshChildView.begin(); child != meshChildView.end(); ++child)
        {
            AZ::SceneAPI::DataTypes::IMeshVertexTangentData* data = azrtti_cast<AZ::SceneAPI::DataTypes::IMeshVertexTangentData*>(child->second.get());
            if (data && setIndex == data->GetTangentSetIndex())
            {
                return data;
            }
        }

        return nullptr;
    }

    bool TangentGenerateComponent::CreateTangentLayer(AZ::SceneAPI::Containers::SceneManifest& manifest,
        const AZ::SceneAPI::Containers::SceneGraph::NodeIndex& nodeIndex,
        size_t numVerts,
        size_t uvSetIndex,
        AZ::SceneAPI::DataTypes::TangentSpace tangentSpace,
        AZ::SceneAPI::Containers::SceneGraph& graph,
        AZ::SceneAPI::DataTypes::IMeshVertexTangentData** outTangentData)
    {
        *outTangentData = nullptr;

        AZStd::shared_ptr<SceneData::GraphData::MeshVertexTangentData> tangentData = AZStd::make_shared<AZ::SceneData::GraphData::MeshVertexTangentData>();
        tangentData->Resize(numVerts);

        AZ_Assert(tangentData, "Failed to allocate tangent data for scene graph.");
        if (!tangentData)
        {
            AZ_Error(AZ::SceneAPI::Utilities::ErrorWindow, false, "Failed to allocate tangent data.\n");
            return false;
        }

        tangentData->SetTangentSetIndex(uvSetIndex);
        tangentData->SetTangentSpace(tangentSpace);

        const AZStd::string tangentGeneratedName = AZStd::string::format("TangentSet_%zu", uvSetIndex);
        const AZStd::string tangentSetName = AZ::SceneAPI::DataTypes::Utilities::CreateUniqueName<SceneData::GraphData::MeshVertexBitangentData>(tangentGeneratedName, manifest);
        AZ::SceneAPI::Containers::SceneGraph::NodeIndex newIndex = graph.AddChild(nodeIndex, tangentSetName.c_str(), tangentData);
        AZ_Assert(newIndex.IsValid(), "Failed to create SceneGraph node for tangent attribute.");
        if (!newIndex.IsValid())
        {
            AZ_Error(AZ::SceneAPI::Utilities::ErrorWindow, false, "Failed to create node in scene graph that stores tangent data.\n");
            return false;
        }
        graph.MakeEndPoint(newIndex);

        *outTangentData = tangentData.get();
        return true;
    }

    AZ::SceneAPI::DataTypes::IMeshVertexBitangentData* TangentGenerateComponent::FindBitangentData(AZ::SceneAPI::Containers::SceneGraph& graph, const AZ::SceneAPI::Containers::SceneGraph::NodeIndex& nodeIndex, AZ::u64 setIndex) const
    {
        const auto nameContentView = AZ::SceneAPI::Containers::Views::MakePairView(graph.GetNameStorage(), graph.GetContentStorage());

        auto meshChildView = AZ::SceneAPI::Containers::Views::MakeSceneGraphChildView<AZ::SceneAPI::Containers::Views::AcceptEndPointsOnly>(graph, nodeIndex, nameContentView.begin(), true);
        for (auto child = meshChildView.begin(); child != meshChildView.end(); ++child)
        {
            AZ::SceneAPI::DataTypes::IMeshVertexBitangentData* data = azrtti_cast<AZ::SceneAPI::DataTypes::IMeshVertexBitangentData*>(child->second.get());
            if (data && setIndex == data->GetBitangentSetIndex())
            {
                return data;
            }
        }

        return nullptr;
    }

    bool TangentGenerateComponent::CreateBitangentLayer(AZ::SceneAPI::Containers::SceneManifest& manifest,
        const AZ::SceneAPI::Containers::SceneGraph::NodeIndex& nodeIndex,
        size_t numVerts,
        size_t uvSetIndex,
        AZ::SceneAPI::DataTypes::TangentSpace tangentSpace,
        AZ::SceneAPI::Containers::SceneGraph& graph,
        AZ::SceneAPI::DataTypes::IMeshVertexBitangentData** outBitangentData)
    {
        *outBitangentData = nullptr;

        AZStd::shared_ptr<AZ::SceneData::GraphData::MeshVertexBitangentData> bitangentData = AZStd::make_shared<AZ::SceneData::GraphData::MeshVertexBitangentData>();
        bitangentData->Resize(numVerts);

        AZ_Assert(bitangentData, "Failed to allocate bitangent data for scene graph.");
        if (!bitangentData)
        {
            AZ_Error(AZ::SceneAPI::Utilities::ErrorWindow, false, "Failed to allocate bitangent data.\n");
            return false;
        }

        bitangentData->SetBitangentSetIndex(uvSetIndex);
        bitangentData->SetTangentSpace(tangentSpace);

        const AZStd::string bitangentGeneratedName = AZStd::string::format("BitangentSet_%zu", uvSetIndex);
        const AZStd::string bitangentSetName = AZ::SceneAPI::DataTypes::Utilities::CreateUniqueName<SceneData::GraphData::MeshVertexBitangentData>(bitangentGeneratedName, manifest);
        AZ::SceneAPI::Containers::SceneGraph::NodeIndex newIndex = graph.AddChild(nodeIndex, bitangentSetName.c_str(), bitangentData);
        AZ_Assert(newIndex.IsValid(), "Failed to create SceneGraph node for bitangent attribute.");
        if (!newIndex.IsValid())
        {
            AZ_Error(AZ::SceneAPI::Utilities::ErrorWindow, false, "Failed to create node in scene graph that stores bitangent data.\n");
            return false;
        }
        graph.MakeEndPoint(newIndex);

        *outBitangentData = bitangentData.get();
        return true;
    }
} // namespace AZ::SceneGenerationComponents