// Copyright (c) 2022, 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 <OgnExponentDatabase.h>
#include <omni/graph/core/ogn/ComputeHelpers.h>
#include <carb/logging/Log.h>
#include <math.h>

namespace omni {
namespace graph {
namespace nodes {

// unnamed namespace to avoid multiple declaration when linking
namespace {
//
bool tryComputeAssumingScalarHalf(OgnExponentDatabase& db)
{
    auto functor = [](auto const& base, auto const& exp, auto& result)
    {
        result = static_cast<pxr::GfHalf>(static_cast<float>(std::pow(static_cast<double>(static_cast<float>(base)), exp)));
    };
    return ogn::compute::tryComputeWithArrayBroadcasting<pxr::GfHalf, int, pxr::GfHalf>(db.inputs.base(),
                                                                                        db.inputs.exponent(),
                                                                                        db.outputs.result(),
                                                                                        functor);
}

template <typename T, typename M>
bool tryComputeAssumingScalarType(OgnExponentDatabase& db)
{
    auto functor = [](auto const& base, auto const& exp, auto& result)
    {
        result = static_cast<M>(std::pow(base, exp));
    };
    return ogn::compute::tryComputeWithArrayBroadcasting<T, int, M>(db.inputs.base(), db.inputs.exponent(),
                                                                    db.outputs.result(), functor);
}

template <size_t N>
bool tryComputeAssumingTupleHalf(OgnExponentDatabase& db)
{
    auto functor = [](auto const& base, auto const& exp, auto& result)
    {
        result = static_cast<pxr::GfHalf>(static_cast<float>(std::pow(static_cast<double>(static_cast<float>(base)), exp)));
    };
    return ogn::compute::tryComputeWithTupleBroadcasting<N, pxr::GfHalf, int, pxr::GfHalf>(db.inputs.base(),
                                                                                           db.inputs.exponent(),
                                                                                           db.outputs.result(),
                                                                                           functor);
}

template <typename T, size_t N, typename M>
bool tryComputeAssumingTupleType(OgnExponentDatabase& db)
{
    auto functor = [](auto const& base, auto const& exp, auto& result)
    {
        result = static_cast<M>(std::pow(base, exp));
    };
    return ogn::compute::tryComputeWithTupleBroadcasting<N, T, int, M>(db.inputs.base(), db.inputs.exponent(),
                                                                       db.outputs.result(), functor);
}
} // unnamed namespace

class OgnExponent
{
public:
    static bool compute(OgnExponentDatabase& db)
    {
        try
        {
            const auto& bType = db.inputs.base().type();
            switch (bType.componentCount)
            {
            case 1:
                switch (bType.baseType)
                {
                case BaseDataType::eInt:
                    return tryComputeAssumingScalarType<int32_t, double>(db);
                case BaseDataType::eInt64:
                    return tryComputeAssumingScalarType<int64_t, double>(db);
                case BaseDataType::eUInt:
                    return tryComputeAssumingScalarType<uint32_t, double>(db);
                case BaseDataType::eUInt64:
                    return tryComputeAssumingScalarType<uint64_t, double>(db);
                case BaseDataType::eUChar:
                    return tryComputeAssumingScalarType<unsigned char, double>(db);
                case BaseDataType::eHalf:
                    return tryComputeAssumingScalarHalf(db);
                case BaseDataType::eDouble:
                    return tryComputeAssumingScalarType<double, double>(db);
                case BaseDataType::eFloat:
                    return tryComputeAssumingScalarType<float, float>(db);
                default:
                    break;
                }
            case 2:
                switch (bType.baseType)
                {
                    case BaseDataType::eInt:
                        return tryComputeAssumingTupleType<int32_t, 2, double>(db);
                    case BaseDataType::eDouble:
                        return tryComputeAssumingTupleType<double, 2, double>(db);
                    case BaseDataType::eFloat:
                        return tryComputeAssumingTupleType<float, 2, float>(db);
                    case BaseDataType::eHalf:
                        return tryComputeAssumingTupleHalf<2>(db);
                    default:
                        break;
                }
            case 3:
                switch (bType.baseType)
                {
                    case BaseDataType::eInt:
                        return tryComputeAssumingTupleType<int32_t, 3, double>(db);
                    case BaseDataType::eDouble:
                        return tryComputeAssumingTupleType<double, 3, double>(db);
                    case BaseDataType::eFloat:
                        return tryComputeAssumingTupleType<float, 3, float>(db);
                    case BaseDataType::eHalf:
                        return tryComputeAssumingTupleHalf<3>(db);
                    default:
                        break;
                }
            case 4:
                switch (bType.baseType)
                {
                    case BaseDataType::eInt:
                        return tryComputeAssumingTupleType<int32_t, 4, double>(db);
                    case BaseDataType::eDouble:
                        return tryComputeAssumingTupleType<double, 4, double>(db);
                    case BaseDataType::eFloat:
                        return tryComputeAssumingTupleType<float, 4, float>(db);
                    case BaseDataType::eHalf:
                        return tryComputeAssumingTupleHalf<4>(db);
                    default:
                        break;
                }
            case 9:
                if (bType.baseType == BaseDataType::eDouble )
                {
                    return tryComputeAssumingTupleType<double, 9, double>(db);
                }
            case 16:
                if (bType.baseType == BaseDataType::eDouble )
                {
                    return tryComputeAssumingTupleType<double, 16, double>(db);
                }
            }
            throw ogn::compute::InputError("Failed to resolve input types");
        }
        catch (ogn::compute::InputError &error)
        {
            db.logError(error.what());
        }
        return false;
    }

    static void onConnectionTypeResolve(const NodeObj& node)
    {
        auto base = node.iNode->getAttributeByToken(node, inputs::base.token());
        auto result = node.iNode->getAttributeByToken(node, outputs::result.token());

        auto bType = base.iAttribute->getResolvedType(base);
        Type newType(BaseDataType::eDouble, bType.componentCount, bType.arrayDepth, bType.role);

        if (bType.baseType != BaseDataType::eUnknown)
        {
            switch (bType.baseType)
            {
            case BaseDataType::eUChar:
            case BaseDataType::eInt:
            case BaseDataType::eUInt:
            case BaseDataType::eInt64:
            case BaseDataType::eUInt64:
                result.iAttribute->setResolvedType(result, newType);
                break;
            default:
                std::array<AttributeObj, 2> attrs { base, result};
                node.iNode->resolveCoupledAttributes(node, attrs.data(), attrs.size());
                break;
            }
        }
    }
};

REGISTER_OGN_NODE()

} // namespace nodes
} // namespace graph
} // namespace omni
