// Copyright (c) 2021-2023, 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 <OgnMultiplyDatabase.h>
#include <omni/graph/core/ogn/ComputeHelpers.h>
#include <carb/logging/Log.h>

namespace omni
{
namespace graph
{
namespace nodes
{

namespace
{
template<typename T>
bool tryComputeAssumingType(OgnMultiplyDatabase& db, size_t count)
{
    auto const& dynamicInputs = db.getDynamicInputs();
    if (dynamicInputs.empty())
    {
        auto functor = [](T const& a, T const& b, T& result)
        {
            result = a * b;
        };
        return ogn::compute::tryComputeWithArrayBroadcasting<T>(db.inputs.a(), db.inputs.b(), db.outputs.product(), functor, count);
    }
    else
    {
        std::vector<ogn::InputAttribute> inputArray{ db.inputs.a(), db.inputs.b() };
        inputArray.reserve(dynamicInputs.size() + 2);
        for (auto const& input : dynamicInputs)
            inputArray.emplace_back(input());

        auto functor = [](const auto& input, auto& result)
        {
            result = result * input;
        };
        return ogn::compute::tryComputeInputsWithArrayBroadcasting<T>(inputArray, db.outputs.product(), functor, count);
    }

}
template<typename T, size_t N>
bool tryComputeAssumingType(OgnMultiplyDatabase& db, size_t count)
{
    auto const& dynamicInputs = db.getDynamicInputs();
    if (dynamicInputs.empty())
    {
        auto functor = [](T const& a, T const& b, T& result)
        {
            result = a * b;
        };
        return ogn::compute::tryComputeWithTupleBroadcasting<N, T>(db.inputs.a(), db.inputs.b(), db.outputs.product(), functor, count);
    }
    else
    {
        std::vector<ogn::InputAttribute> inputArray{ db.inputs.a(), db.inputs.b() };
        inputArray.reserve(dynamicInputs.size() + 2);
        for (auto const& input : dynamicInputs)
            inputArray.emplace_back(input());

        auto functor = [](const auto& input, auto& result)
        {
            result = result * input;
        };
        return ogn::compute::tryComputeInputsWithTupleBroadcasting<N, T>(inputArray, db.outputs.product(), functor, count);
    }
}
} // namespace

class OgnMultiply
{
public:
    static size_t computeVectorized(OgnMultiplyDatabase& db, size_t count)
    {
        auto& productType = db.outputs.product().type();
        try
        {
            switch (productType.baseType)
            {
            case BaseDataType::eDouble:
                switch (productType.componentCount)
                {
                    case 1:  return tryComputeAssumingType<double>(db, count);
                    case 2:  return tryComputeAssumingType<double, 2>(db, count);
                    case 3:  return tryComputeAssumingType<double, 3>(db, count);
                    case 4:  return tryComputeAssumingType<double, 4>(db, count);
                    case 9:  return tryComputeAssumingType<double, 9>(db, count);
                    case 16: return tryComputeAssumingType<double, 16>(db, count);
                    default: break;
                }
            case BaseDataType::eFloat:
                switch (productType.componentCount)
                {
                    case 1: return tryComputeAssumingType<float>(db, count);
                    case 2: return tryComputeAssumingType<float, 2>(db, count);
                    case 3: return tryComputeAssumingType<float, 3>(db, count);
                    case 4: return tryComputeAssumingType<float, 4>(db, count);
                    default: break;
                }
            case BaseDataType::eHalf:
                switch (productType.componentCount)
                {
                    case 1: return tryComputeAssumingType<pxr::GfHalf>(db, count);
                    case 2: return tryComputeAssumingType<pxr::GfHalf, 2>(db, count);
                    case 3: return tryComputeAssumingType<pxr::GfHalf, 3>(db, count);
                    case 4: return tryComputeAssumingType<pxr::GfHalf, 4>(db, count);
                    default: break;
                }
            case BaseDataType::eInt:
                switch (productType.componentCount)
                {
                    case 1: return tryComputeAssumingType<int32_t>(db, count);
                    case 2: return tryComputeAssumingType<int32_t, 2>(db, count);
                    case 3: return tryComputeAssumingType<int32_t, 3>(db, count);
                    case 4: return tryComputeAssumingType<int32_t, 4>(db, count);
                    default: break;
                }
                ;
            case BaseDataType::eInt64:
                return tryComputeAssumingType<int64_t>(db, count);
            case BaseDataType::eUChar:
                return tryComputeAssumingType<unsigned char>(db, count);
            case BaseDataType::eUInt:
                return tryComputeAssumingType<uint32_t>(db, count);
            case BaseDataType::eUInt64:
                return tryComputeAssumingType<uint64_t>(db, count);
            default: break;
            }

            db.logWarning("Failed to resolve input types");
        }
        catch (ogn::compute::InputError &error)
        {
            db.logWarning("OgnMultiply: %s", error.what());
        }
        return 0;
    }

    static void onConnectionTypeResolve(const NodeObj& node)
    {
        auto totalCount = node.iNode->getAttributeCount(node);
        std::vector<AttributeObj> allAttributes(totalCount);
        node.iNode->getAttributes(node, allAttributes.data(), totalCount);

        std::vector<AttributeObj> attributes;
        std::vector<uint8_t> componentCounts;
        std::vector<uint8_t> arrayDepths;
        std::vector<AttributeRole> roles;

        attributes.reserve(totalCount - 2);
        componentCounts.reserve(totalCount - 2);
        arrayDepths.reserve(totalCount - 2);
        roles.reserve(totalCount - 2);

        uint8_t maxArrayDepth = 0;
        uint8_t maxComponentCount = 0;

        for (auto const& attr : allAttributes)
        {
            if (attr.iAttribute->getPortType(attr) == AttributePortType::kAttributePortType_Input)
            {
                auto resolvedType = attr.iAttribute->getResolvedType(attr);

                // if some inputs are not connected stop - the output port resolution is only completed when all inputs
                // are connected
                if (resolvedType.baseType == BaseDataType::eUnknown)
                    return;

                componentCounts.push_back(resolvedType.componentCount);
                arrayDepths.push_back(resolvedType.arrayDepth);
                roles.push_back(resolvedType.role);
                maxComponentCount = std::max(maxComponentCount, resolvedType.componentCount);
                maxArrayDepth = std::max(maxArrayDepth, resolvedType.arrayDepth);

                attributes.push_back(attr);
            }
        }

        auto output = node.iNode->getAttributeByToken(node, outputs::product.token());
        attributes.push_back(output);
        // All inputs and the output should have the same tuple count
        componentCounts.push_back(maxComponentCount);
        // Allow for a mix of singular and array inputs. If any input is an array, the output must be an array
        arrayDepths.push_back(maxArrayDepth);
        // Copy the attribute role from the resolved type to the output type
        roles.push_back(AttributeRole::eUnknown);

        node.iNode->resolvePartiallyCoupledAttributes(
            node, attributes.data(), componentCounts.data(), arrayDepths.data(), roles.data(), attributes.size());
    }
};

REGISTER_OGN_NODE()

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