// 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 <OgnTutorialBundleDataDatabase.h>

using omni::graph::core::Type;
using omni::graph::core::NameToken;
using omni::graph::core::BaseDataType;


namespace
{
// The function of this node is to double the values for all bundled attributes
// with integral types, including tuples and arrays.
//
// The parameter is a ogn::RuntimeAttribute<kOgnOutput, ogn::kCpu>, which is the type of data returned when iterating
// over an output bundle on the CPU.
// It contains a description of the attribute within the bundle and access to the attribute's data.
// BundledInput is a similar type, which is what you get when iterating over an input bundle. The main difference
// between the two is the ability to modify the attribute or its data.
template <typename POD>
bool doubleSimple(ogn::RuntimeAttribute<ogn::kOgnOutput, ogn::kCpu>& bundledAttribute)
{
    // When an attribute is cast to the wrong type (e.g. an integer attribute is extracted with a float
    // template parameter on the get<,>() method) a nullptr is returned. That can be used to determine
    // the attribute type. You can also use the bundledAttribute.type() method to access the full type
    // information and select a code path using that.
    auto podValue = bundledAttribute.get<POD>();
    if (podValue)
    {
        *podValue *= 2;
        return true;
    }
    return false;
};
// Array and tuple data has iterator capabilities for easy access to individual elements
template <typename CppType>
bool doubleArray(ogn::RuntimeAttribute<ogn::kOgnOutput, ogn::kCpu>& bundledAttribute)
{
    // Strings and paths look like uint8_t[] but are not, so don't process them
    if ((bundledAttribute.type().role == AttributeRole::eText) || (bundledAttribute.type().role == AttributeRole::ePath))
    {
        return false;
    }
    auto arrayValue = bundledAttribute.get<CppType>();
    if (arrayValue)
    {
        for (auto& arrayElement : *arrayValue)
        {
            arrayElement *= 2;
        }
        return true;
    }
    return false;
};
// Tuple arrays must have nested iteration
template <typename CppType>
bool doubleTupleArray(ogn::RuntimeAttribute<ogn::kOgnOutput, ogn::kCpu>& bundledAttribute)
{
    auto tupleSize = bundledAttribute.type().componentCount;
    auto tupleArrayValue = bundledAttribute.get<CppType>();
    if (tupleArrayValue)
    {
        for (auto& arrayElement : *tupleArrayValue)
        {
            for (size_t i=0; i<tupleSize; ++i)
            {
                arrayElement[i] *= 2;
            }
        }
        return true;
    }
    return false;
};
}

class OgnTutorialBundleData
{
public:
    static bool compute(OgnTutorialBundleDataDatabase& db)
    {
        // Bundle attributes are extracted from the database in the same way as any other attribute.
        // The only difference is that a different interface class is provided, suited to bundle manipulation.
        const auto& inputBundle = db.inputs.bundle();
        auto& outputBundle = db.outputs.bundle();

        for (const auto& bundledAttribute : inputBundle)
        {
            auto podValue = bundledAttribute.get<int>();
        }

        // Copying the entire bundle is more efficient than adding one member at a time. As the output bundle
        // will have the same attributes as the input, even though the values will be different, this is the best
        // approach. If the attribute lists were different then you would copy or create the individual attributes
        // as required.
        outputBundle = inputBundle;

        // Now walk the bundle to look for types to be modified
        for (auto& bundledAttribute : outputBundle)
        {
            // This shows how using a templated function can simplify handling of several different bundled
            // attribute types. The data types for each of the POD attributes is fully explained in the documentation
            // page titled "Attribute Data Types". The list of available POD data types is:
            //
            //      bool
            //      double
            //      float
            //      pxr::GfHalf
            //      int
            //      int64_t
            //      NameToken
            //      uint8_t
            //      uint32_t
            //      uint64_t
            //
            if (doubleSimple<int>(bundledAttribute)) continue;
            if (doubleSimple<int64_t>(bundledAttribute)) continue;
            if (doubleSimple<uint8_t>(bundledAttribute)) continue;
            if (doubleSimple<uint32_t>(bundledAttribute)) continue;
            if (doubleSimple<uint64_t>(bundledAttribute)) continue;

            // Plain ints are the only integral types supporting tuples. Double those here
            if (doubleArray<int[2]>(bundledAttribute)) continue;
            if (doubleArray<int[3]>(bundledAttribute)) continue;
            if (doubleArray<int[4]>(bundledAttribute)) continue;

            // Arrays are looped differently than tupls so they are also handled differently
            if (doubleArray<int[]>(bundledAttribute)) continue;
            if (doubleArray<int64_t[]>(bundledAttribute)) continue;
            if (doubleArray<uint8_t[]>(bundledAttribute)) continue;
            if (doubleArray<uint32_t[]>(bundledAttribute)) continue;
            if (doubleArray<uint64_t[]>(bundledAttribute)) continue;

            // Tuple arrays require a two level iteration
            if (doubleTupleArray<int[][2]>(bundledAttribute)) continue;
            if (doubleTupleArray<int[][3]>(bundledAttribute)) continue;
            if (doubleTupleArray<int[][4]>(bundledAttribute)) continue;

            // Any attributes not passing the above tests are not integral and therefore not handled by this node.
        }
        return true;
    }
};

REGISTER_OGN_NODE()
