// Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved.
//
// NVIDIA CORPORATION and its licensors retain all intellectual property
// and proprietary rights in and to this software, related documentation
// and any modifications thereto.  Any use, reproduction, disclosure or
// distribution of this software and related documentation without an express
// license agreement from NVIDIA CORPORATION is strictly prohibited.
//
#include <OgnDistance3DDatabase.h>
#include <omni/graph/core/ogn/ComputeHelpers.h>
#include <carb/logging/Log.h>

#include <cmath>

namespace omni
{
namespace graph
{
namespace nodes
{

namespace
{
template<typename T>
bool tryComputeAssumingType(OgnDistance3DDatabase& db, size_t count)
{
    auto functor = [](auto const& a, auto const& b, auto& result)
    {
        result = sqrt((b[0] - a[0]) * (b[0] - a[0]) + (b[1] - a[1]) * (b[1] - a[1]) + (b[2] - a[2]) * (b[2] - a[2]));
    };
    return ogn::compute::tryComputeWithArrayBroadcasting<T[3], T[3], T>(db.inputs.a(), db.inputs.b(), db.outputs.distance(), functor, count);
}
} // namespace

class OgnDistance3D
{
public:
    static bool computeVectorized(OgnDistance3DDatabase& db, size_t count)
    {
        try
        {
            if (tryComputeAssumingType<double>(db, count)) return true;
            else if (tryComputeAssumingType<float>(db, count)) return true;
            else if (tryComputeAssumingType<pxr::GfHalf>(db, count)) return true;
            else
            {
                db.logWarning("OgnDistance3D: Failed to resolve input types");
            }
        }
        catch (ogn::compute::InputError &error)
        {
            db.logWarning("OgnDistance3D: %s", error.what());
        }
        return false;
    }

    static void onConnectionTypeResolve(const NodeObj& node)
    {
        auto a = node.iNode->getAttributeByToken(node, inputs::a.token());
        auto b = node.iNode->getAttributeByToken(node, inputs::b.token());
        auto distance = node.iNode->getAttributeByToken(node, outputs::distance.token());

        auto typeA = a.iAttribute->getResolvedType(a);
        auto typeB = a.iAttribute->getResolvedType(b);

        const Type invalid;
        // Require a and b to be resolved before determining result's type
        if (typeA != invalid && typeB != invalid)
        {
            uint8_t ad = std::max(typeA.arrayDepth, typeB.arrayDepth);
            AttributeObj attrs[]  = { a, b, distance };
            uint8_t tupleCounts[] = { 3, 3, 1 };
            uint8_t arrayDepths[] = { typeA.arrayDepth, typeB.arrayDepth, ad };
            AttributeRole roles[] = { AttributeRole::eNone, AttributeRole::eNone, AttributeRole::eNone };

            node.iNode->resolvePartiallyCoupledAttributes(node, attrs, tupleCounts, arrayDepths, roles, 3);
        }
    }
};

REGISTER_OGN_NODE()

} // namespace nodes
} // namespace graph
} // namespace omni
