// 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 <OgnTutorialCpuGpuBundlesDatabase.h>
extern "C" void cpuGpuDotProductCPU(float const(*p1)[3], float const(*p2)[3], float*, size_t);
extern "C" void cpuGpuDotProductGPU(float const (**p1)[3], float const (**p2)[3], float**, size_t);

namespace omni {
namespace graph {
namespace tutorials {

class OgnTutorialCpuGpuBundles
{
public:
    static bool compute(OgnTutorialCpuGpuBundlesDatabase& db)
    {
        const auto& gpu = db.inputs.gpu();
        // Bundles are merely abstract representations of a collection of attributes so you don't have to do anything
        // different when they are marked for GPU, or ANY memory location.
        const auto& cpuBundle = db.inputs.cpuBundle();
        const auto& gpuBundle = db.inputs.gpuBundle();
        auto& outputBundle = db.outputs.cpuGpuBundle();

        // Assign the correct destination bundle to the output based on the gpu flag
        if (gpu)
        {
            outputBundle = gpuBundle;
        }
        else
        {
            outputBundle = cpuBundle;
        }

        // Get the attribute references. They're the same whether the bundles are on the CPU or GPU
        const auto pointsCpuAttribute = cpuBundle.attributeByName(db.tokens.points);
        const auto pointsGpuAttribute = gpuBundle.attributeByName(db.tokens.points);
        auto dotProductAttribute = outputBundle.attributeByName(db.tokens.dotProducts);
        if (! dotProductAttribute.isValid())
        {
            dotProductAttribute = outputBundle.addAttribute(db.tokens.dotProducts, Type(BaseDataType::eFloat, 1, 1));
        }

        // Find the bundle contents to be processed
        if (gpu)
        {
            const auto points1 = pointsCpuAttribute.getGpu<float[][3]>();
            const auto points2 = pointsGpuAttribute.get<float[][3]>();
            auto dotProducts = dotProductAttribute.getGpu<float[]>();
            if (! points1)
            {
                db.logWarning("Skipping compute - No valid float[3][] attribute named '%s' on the CPU bundle", db.tokenToString(db.tokens.points));
                return false;
            }
            if (! points2)
            {
                db.logWarning("Skipping compute - No valid float[3][] attribute named '%s' on the GPU bundle", db.tokenToString(db.tokens.points));
                return false;
            }
            if (points1.size() != points2.size())
            {
                db.logWarning("Skipping compute - Point arrays are different sizes (%zu and %zu)", points1.size(), points2.size());
                return false;
            }
            dotProducts.resize(points1.size());
            if (! dotProducts)
            {
                db.logWarning("Skipping compute - No valid float[] attribute named '%s' on the output bundle", db.tokenToString(db.tokens.dotProducts));
                return false;
            }
            cpuGpuDotProductGPU(points1(), points2(), dotProducts(), points1.size());
        }
        else
        {
            const auto points1 = pointsCpuAttribute.get<float[][3]>();
            const auto points2 = pointsGpuAttribute.getCpu<float[][3]>();
            auto dotProducts = dotProductAttribute.getCpu<float[]>();
            if (! points1)
            {
                db.logWarning("Skipping compute - No valid float[3][] attribute named '%s' on the CPU bundle", db.tokenToString(db.tokens.points));
                return false;
            }
            if (! points2)
            {
                db.logWarning("Skipping compute - No valid float[3][] attribute named '%s' on the GPU bundle", db.tokenToString(db.tokens.points));
                return false;
            }
            if (points1.size() != points2.size())
            {
                db.logWarning("Skipping compute - Point arrays are different sizes (%zu and %zu)", points1.size(), points2.size());
                return false;
            }
            dotProducts.resize(points1.size());
            if (! dotProducts)
            {
                db.logWarning("Skipping compute - No valid dot product attribute on the output bundle");
                return false;
            }
            cpuGpuDotProductCPU(points1->data(), points2->data(), dotProducts->data(), points1.size());
        }
        return true;
    }
};

REGISTER_OGN_NODE()

} // namespace tutorials
} // namespace graph
} // namespace omni
