// 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 <carb/logging/Log.h>

#include <omni/graph/core/ogn/ComputeHelpers.h>

#include <OgnAbsoluteDatabase.h>
#include <NumericUtils.h>
#include <cmath>
#include <type_traits>

namespace omni
{
namespace graph
{
namespace nodes
{

namespace
{

template <typename T>
T ogAbs(T const& input) noexcept
{
    if constexpr (std::is_unsigned_v<T>)
    {
        return input;
    }
    else
    {
        return std::abs(input);
    }
}

template <typename T, size_t N>
struct ComputeAbsoluteValueAssumingType
{
    bool operator()(OgnAbsoluteDatabase& db) const
    {
        if constexpr (N == 1)
        {
            auto functor = [](T const& input, T& absolute) { absolute = ogAbs(input); };
            return ogn::compute::tryComputeWithArrayBroadcasting<T, T>(db.inputs.input(), db.outputs.absolute(), functor);
        }
        else
        {
            auto functor = [](auto const& input, auto& absolute)
            {
                for (size_t i = 0; i < N; ++i)
                {
                    absolute[i] = ogAbs(input[i]);
                }
            };
            return ogn::compute::tryComputeWithArrayBroadcasting<T[N], T[N]>(
                db.inputs.input(), db.outputs.absolute(), functor);
        }
    }
};

}

class OgnAbsolute
{
public:
    static bool compute(OgnAbsoluteDatabase& db)
    {
        try
        {
            auto const& type = db.inputs.input().type();
            if (!callForNumericAttribute<ComputeAbsoluteValueAssumingType>(db, type))
            {
                throw ogn::compute::InputError("Failed to resolve input type");
            }
        }
        catch (ogn::compute::InputError& error)
        {
            db.logError("OgnAbsolute: %s", error.what());
        }
        return false;
    }

    static void onConnectionTypeResolve(const NodeObj& node)
    {
        auto input = node.iNode->getAttributeByToken(node, inputs::input.token());
        auto absolute = node.iNode->getAttributeByToken(node, outputs::absolute.token());
        auto inputType = input.iAttribute->getResolvedType(input);

        if (inputType.baseType != BaseDataType::eUnknown)
        {
            absolute.iAttribute->setResolvedType(absolute, inputType);
        }
    }
};

REGISTER_OGN_NODE()

}
}
}
