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

// The method used by the computation to perform the passthrough
// 0 = simple copy, no vectorization
#define TUTO_ABI_PASSTHROUGH_METHOD_SIMPLE 0

// 1 = vectorized copy by indexing the instance directly per attribute
#define TUTO_ABI_PASSTHROUGH_METHOD_ATTR 1

// 2 = vectorized copy by mutating attribute data handles
#define TUTO_ABI_PASSTHROUGH_METHOD_MUTATE 2

// 3 = vectorized copy using raw data
#define TUTO_ABI_PASSTHROUGH_METHOD_RAW 3

//By default, use the most efficient method
#define TUTO_ABI_PASSTHROUGH_METHOD TUTO_ABI_PASSTHROUGH_METHOD_RAW

// This node perform a copy of its input to its output
class OgnTutorialVectorizedABIPassthrough
{
public:
#if TUTO_ABI_PASSTHROUGH_METHOD == TUTO_ABI_PASSTHROUGH_METHOD_SIMPLE
    // begin-regular
    static bool compute(GraphContextObj const& contextObj, NodeObj const& nodeObj)
    {
        NodeContextHandle nodeHandle = nodeObj.nodeContextHandle;

        auto inputValueAttr = 
            getAttributeR(contextObj, nodeHandle, Token("inputs:value"), kAccordingToContextIndex);
        const float* inputValue = getDataR<float>(contextObj, inputValueAttr);

        auto outputValueAttr =
            getAttributeW(contextObj, nodeHandle, Token("outputs:value"), kAccordingToContextIndex);
        float* outputValue = getDataW<float>(contextObj, outputValueAttr);

        if(inputValue && outputValue)
        {
            *outputValue = *inputValue;
            return true;
        }

        return false;
    }
    // end-regular
#elif TUTO_ABI_PASSTHROUGH_METHOD == TUTO_ABI_PASSTHROUGH_METHOD_ATTR
    // begin-attr
    static size_t computeVectorized(GraphContextObj const& contextObj, NodeObj const& nodeObj, size_t count)
    {
        GraphContextObj const* contexts = nullptr;
        NodeObj const* nodes = nullptr;

        //When using auto instancing, similiar graphs can get merged together, and computed vectorized
        // In such case, each instance represent a different node in a different graph
        // Accessing the data either through the provided node, or through the actual auto-insance node would work properly
        // But any other ABI call requiring the node would need to provide the proper node
        // While not necessary in this contezt, do the work of using the proper auto-instance node
        // in order to demonstrate how to use it.
        size_t handleCount = nodeObj.iNode->getAutoInstances(nodeObj, contexts, nodes);

        auto nodeHandle = [&](InstanceIndex index) -> NodeContextHandle
        { return nodes[handleCount == 1 ? 0 : index.index].nodeContextHandle; };

        auto context = [&](InstanceIndex index) -> GraphContextObj
        { return contexts[handleCount == 1 ? 0 : index.index]; };

        size_t ret = 0;
        const float* inputValue{ nullptr };
        float* outputValue{ nullptr };
        auto inToken = Token("inputs:value");
        auto outToken = Token("outputs:value");

        for (InstanceIndex idx{ 0 }; idx < InstanceIndex{ count }; ++idx)
        {
            auto inputValueAttr = getAttributeR(context(idx), nodeHandle(idx), inToken, idx);
            inputValue = getDataR<float>(context(idx), inputValueAttr);

            auto outputValueAttr = getAttributeW(context(idx), nodeHandle(idx), outToken, idx);
            outputValue = getDataW<float>(context(idx), outputValueAttr);

            if(inputValue && outputValue)
            {
                *outputValue = *inputValue;
                ++ret;
            }
        }

        return ret;
    }
    // end-attr
#elif TUTO_ABI_PASSTHROUGH_METHOD == TUTO_ABI_PASSTHROUGH_METHOD_MUTATE
    // begin-mutate
    static size_t computeVectorized(GraphContextObj const& contextObj, NodeObj const& nodeObj, size_t count)
    {
        NodeContextHandle nodeHandle = nodeObj.nodeContextHandle;

        size_t ret = 0;
        const float* inputValue{ nullptr };
        float* outputValue{ nullptr };
        auto inputValueAttr = 
            getAttributeR(contextObj, nodeHandle, Token("inputs:value"), kAccordingToContextIndex);
        auto outputValueAttr =
            getAttributeW(contextObj, nodeHandle, Token("outputs:value"), kAccordingToContextIndex);

        while (count--)
        {
            inputValue = getDataR<float>(contextObj, inputValueAttr);
            outputValue = getDataW<float>(contextObj, outputValueAttr);

            if(inputValue && outputValue)
            {
                *outputValue = *inputValue;
                ++ret;
            }
            inputValueAttr = contextObj.iAttributeData->moveToAnotherInstanceR(contextObj, inputValueAttr, 1);
            outputValueAttr = contextObj.iAttributeData->moveToAnotherInstanceW(contextObj, outputValueAttr, 1);
        }

        return ret;
    }
    // end-mutate
#elif TUTO_ABI_PASSTHROUGH_METHOD == TUTO_ABI_PASSTHROUGH_METHOD_RAW
    // begin-raw
    static size_t computeVectorized(GraphContextObj const& contextObj, NodeObj const& nodeObj, size_t count)
    {
        NodeContextHandle nodeHandle = nodeObj.nodeContextHandle;

        auto inputValueAttr = 
            getAttributeR(contextObj, nodeHandle, Token("inputs:value"), kAccordingToContextIndex);
        const float* inputValue = getDataR<float>(contextObj, inputValueAttr);

        auto outputValueAttr =
            getAttributeW(contextObj, nodeHandle, Token("outputs:value"), kAccordingToContextIndex);
        float* outputValue = getDataW<float>(contextObj, outputValueAttr);

        if(inputValue && outputValue)
        {
            memcpy(outputValue, inputValue, count * sizeof(float));
            return count;
        }

        return 0;
    }
    // end-raw
#endif
};

REGISTER_OGN_NODE()
