// Copyright (c) 2022, 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 <OgnDotProductDatabase.h>
#include <omni/graph/core/ogn/ComputeHelpers.h>
#include <omni/math/linalg/vec.h>
#include <carb/logging/Log.h>
#include <cmath>

namespace omni {
namespace graph {
namespace nodes {

// unnamed namespace to avoid multiple declaration when linking
namespace {

template<typename T, size_t N>
bool tryComputeAssumingType(OgnDotProductDatabase& db)
{
    auto functor = [](auto const& a, auto const& b, auto& product)
    {
        const auto& vecA = *reinterpret_cast<const omni::math::linalg::base_vec<T, N>*>(a);
        const auto& vecB = *reinterpret_cast<const omni::math::linalg::base_vec<T, N>*>(b);
        product = vecA.Dot(vecB);
    };
    return ogn::compute::tryComputeWithArrayBroadcasting<T[N], T[N], T>(db.inputs.a(), db.inputs.b(), db.outputs.product(), functor);
}

} // namespace

class OgnDotProduct
{
public:
    static bool compute(OgnDotProductDatabase& db)
    {
        try
        {
            auto& aType = db.inputs.a().type();
            
            switch(aType.baseType)
            {
                case BaseDataType::eDouble:
                    switch(aType.componentCount)
                    {
                        case 2: return tryComputeAssumingType<double, 2>(db);
                        case 3: return tryComputeAssumingType<double, 3>(db);
                        case 4: return tryComputeAssumingType<double, 4>(db);
                    }
                case BaseDataType::eFloat:
                    switch(aType.componentCount)
                    {
                        case 2: return tryComputeAssumingType<float, 2>(db);
                        case 3: return tryComputeAssumingType<float, 3>(db);
                        case 4: return tryComputeAssumingType<float, 4>(db);
                    }
                case BaseDataType::eHalf:
                    switch(aType.componentCount)
                    {
                        case 2: return tryComputeAssumingType<pxr::GfHalf, 2>(db);
                        case 3: return tryComputeAssumingType<pxr::GfHalf, 3>(db);
                        case 4: return tryComputeAssumingType<pxr::GfHalf, 4>(db);
                    }
                default: db.logError("Failed to resolve input types");
            }
        }
        catch (ogn::compute::InputError &error)
        {
            db.logError("%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 product = node.iNode->getAttributeByToken(node, outputs::product.token());
        
        // Require inputs to be resolved before determining product's type
        auto aType = a.iAttribute->getResolvedType(a);
        auto bType = b.iAttribute->getResolvedType(b);
        if (aType.baseType != BaseDataType::eUnknown && bType.baseType != BaseDataType::eUnknown)
        {
            std::array<AttributeObj, 3> attrs { a, b, product };
            // a, b should all have the same tuple count
            std::array<uint8_t, 3> tupleCounts {
                aType.componentCount,
                bType.componentCount,
                1
            };
            std::array<uint8_t, 3> arrayDepths {
                aType.arrayDepth,
                bType.arrayDepth,
                // Allow for a mix of singular and array inputs. If any input is an array, the output must be an array
                std::max(aType.arrayDepth, bType.arrayDepth)
            };
            std::array<AttributeRole, 3> rolesBuf {
                aType.role,
                bType.role,
                AttributeRole::eNone
            };
            node.iNode->resolvePartiallyCoupledAttributes(node, attrs.data(), tupleCounts.data(), arrayDepths.data(), rolesBuf.data(), attrs.size());
        }
    }
};

REGISTER_OGN_NODE()

} // namespace nodes
} // namespace graph
} // namespace omni
// end-compute-helpers
