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

namespace omni {
namespace graph {
namespace tutorials {

// Only the pointf[3][] type is accepted. Make a shortcut to the type information describing it for comparison
core::Type acceptedType(core::BaseDataType::eFloat, 3, 1, core::AttributeRole::ePosition);

// Template for reporting type inconsistencies. The data types are different for the attributes but the checks are
// the same so this avoids duplication
template <typename P1, typename P2, typename P3>
bool verifyDataTypes(OgnTutorialCpuGpuExtendedDatabase& db, const P1& points1, const P2& points2, P3& sums, const char* type)
{
    if (! points1)
    {
        db.logWarning("Skipping compute - The %s attribute was not a valid pointf[3][]", type);
    }
    else if (! points2)
    {
        db.logWarning("Skipping compute - The %s attribute was not a valid pointf[3][]", type);
    }
    else if (! sums)
    {
        db.logWarning("Skipping compute - The %s output attribute was not a valid pointf[3][]", type);
    }
    else if (points1.size() != points2.size())
    {
        db.logWarning("Skipping compute - Point arrays are different sizes (%zu and %zu)", points1.size(), points2.size());
    }
    else 
    {
        sums.resize(points1.size());
        return true;
    }
    return false;
}

class OgnTutorialCpuGpuExtended
{
    bool m_allAttributesResolved{ false };
public:
    static bool compute(OgnTutorialCpuGpuExtendedDatabase& db)
    {
        if (! db.sharedState<OgnTutorialCpuGpuExtended>().m_allAttributesResolved)
        {
            db.logWarning("All types are not yet resolved. Cannot run the compute.");
            return false;
        }

        const auto& gpu = db.inputs.gpu();
        const auto cpuData = db.inputs.cpuData();
        const auto gpuData = db.inputs.gpuData();
        auto cpuGpuSum = db.outputs.cpuGpuSum();

        if ((cpuData.type() != acceptedType) || (gpuData.type() != acceptedType) || (cpuGpuSum.type() != acceptedType))
        {
            db.logWarning("Skipping compute - All of the attributes do not have the accepted resolved type pointf[3][]");
            return false;
        }

        if (gpu)
        {
            // Computation on the GPU has been requested so get the GPU versions of the attribute data
            const auto points1 = cpuData.getGpu<float[][3]>();
            const auto points2 = gpuData.get<float[][3]>();
            auto sums = cpuGpuSum.getGpu<float[][3]>();
            if (!verifyDataTypes(db, points1, points2, sums, "GPU"))
            {
                return false;
            }
            cpuGpuSumGPU(points1(), points2(), sums());
        }
        else
        {
            // Computation on the CPU has been requested so get the CPU versions of the attribute data
            const auto points1 = cpuData.get<float[][3]>();
            const auto points2 = gpuData.getCpu<float[][3]>();
            auto sums = cpuGpuSum.getCpu<float[][3]>();
            if (!verifyDataTypes(db, points1, points2, sums, "CPU"))
            {
                return false;
            }
            cpuGpuSumCPU(points1->data(), points2->data(), sums->data(), points1.size());
        }
        return true;
    }

    static void onConnectionTypeResolve(const NodeObj& nodeObj)
    {
        // If any one type is resolved the others should resolve to the same type. Calling this helper function
        // makes that happen automatically. If it returns false then the resolution failed for some reason. The
        // node's user data, which is just a copy of this class, is used to keep track of the resolution state so
        // that the compute method can quickly exit when the types are not resolved.
        AttributeObj attributes[3] {
            nodeObj.iNode->getAttributeByToken(nodeObj, inputs::cpuData.token()),
            nodeObj.iNode->getAttributeByToken(nodeObj, inputs::gpuData.token()),
            nodeObj.iNode->getAttributeByToken(nodeObj, outputs::cpuGpuSum.token())
        };
        auto& state = OgnTutorialCpuGpuExtendedDatabase::sInternalState<OgnTutorialCpuGpuExtended>(nodeObj);
        state.m_allAttributesResolved = nodeObj.iNode->resolveCoupledAttributes(nodeObj, attributes, 3);
    }
};

REGISTER_OGN_NODE()

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