// Copyright (c) 2019-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 <OgnTestAddAnyTypeAnyMemoryDatabase.h>


namespace omni
{
namespace graph
{
namespace examples
{

extern "C" void testAddAnyTypeAnyMemoryEP(float const*, float const*, float*);

class OgnTestAddAnyTypeAnyMemory
{
public:
    static bool compute(OgnTestAddAnyTypeAnyMemoryDatabase& db)
    {
        //This is a test node not supposed to be used in a real graph, but only for test
        // we assume that it will only be used with float inputs
        auto const& vec = *db.inputs.vec().getCpu<float[3]>();
        auto const& scalar = *db.inputs.scalar().getCpu<float>();
        auto& resCpu = *db.outputs.outCpu().get<float[3]>();
        for (size_t i = 0; i < 3; ++i)
            resCpu[i] = vec[i] + scalar;

        auto const& gpu_vec = *db.inputs.vec().getGpu<float[3]>();
        auto const& gpu_scalar = *db.inputs.scalar().getGpu<float>();
        auto& resGpu = *db.outputs.outGpu().get<float[3]>();
        testAddAnyTypeAnyMemoryEP(gpu_vec, &gpu_scalar, resGpu);

        return true;
    }

    static void onConnectionTypeResolve(const NodeObj& nodeObj)
    {
        AttributeObj attributes[4]{ nodeObj.iNode->getAttributeByToken(nodeObj, inputs::vec.token()),
                                    nodeObj.iNode->getAttributeByToken(nodeObj, inputs::scalar.token()),
                                    nodeObj.iNode->getAttributeByToken(nodeObj, outputs::outCpu.token()),
                                    nodeObj.iNode->getAttributeByToken(nodeObj, outputs::outGpu.token()) };
        uint8_t tupleBuff[4] = {3, 1, 3, 3};
        uint8_t arrayBuff[4] = {0, 0, 0, 0};
        AttributeRole roleBuff[4] = { AttributeRole::eNone, AttributeRole::eNone, AttributeRole::eNone,
                                      AttributeRole::eNone };
        nodeObj.iNode->resolvePartiallyCoupledAttributes(nodeObj, attributes, tupleBuff, arrayBuff, roleBuff, 4);
    }
};

REGISTER_OGN_NODE()

}
}
}
