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

namespace omni {
namespace graph {
namespace nodes {

// unnamed namespace to avoid multiple declaration when linking
namespace {

/**
  _____   ____    _   _  ____ _______    _____ ____  _______     __
 |  __ \ / __ \  | \ | |/ __ \__   __|  / ____/ __ \|  __ \ \   / /
 | |  | | |  | | |  \| | |  | | | |    | |   | |  | | |__) \ \_/ /
 | |  | | |  | | | . ` | |  | | | |    | |   | |  | |  ___/ \   /
 | |__| | |__| | | |\  | |__| | | |    | |___| |__| | |      | |
 |_____/ \____/  |_| \_|\____/  |_|     \_____\____/|_|      |_|

This node uses a large cascading "if" to select operation type, which is not efficient. It will be eventually be
refactored but until then do not propagate this anti-pattern. Thanks for keeping things fast!
*/

/**
 * Used when input type is resolved as non-int numeric type other than Half
 */
template <typename T>
bool tryComputeAssumingType(OgnTrigDatabase& db, NameToken operation, size_t count)
{

    if (operation == db.tokens.SIN) // Sine
    {
        auto functor = [](auto const& a, auto& result) { result = static_cast<T>(std::sin(pxr::GfDegreesToRadians(a))); };
        return ogn::compute::tryCompute(db.inputs.a().template get<T>(), db.outputs.result().template get<T>(), functor, count);
    }
    else if (operation == db.tokens.COS) // Cosine
    {
        auto functor = [](auto const& a, auto& result) { result = static_cast<T>(std::cos(pxr::GfDegreesToRadians(a))); };
        return ogn::compute::tryCompute(db.inputs.a().template get<T>(), db.outputs.result().template get<T>(), functor, count);
    }
    else if (operation == db.tokens.TAN) // Tangent
    {
        auto functor = [](auto const& a, auto& result) { result = static_cast<T>(std::tan(pxr::GfDegreesToRadians(a))); };
        return ogn::compute::tryCompute(db.inputs.a().template get<T>(), db.outputs.result().template get<T>(), functor, count);
    }
    else if (operation == db.tokens.ARCSIN) // Arcsine
    {
        auto functor = [](auto const& a, auto& result)
        { result = static_cast<T>(pxr::GfRadiansToDegrees(std::asin(a))); };
        return ogn::compute::tryCompute(db.inputs.a().template get<T>(), db.outputs.result().template get<T>(), functor, count);
    }
    else if (operation == db.tokens.ARCCOS) // Arccosine
    {
        auto functor = [](auto const& a, auto& result)
        { result = static_cast<T>(pxr::GfRadiansToDegrees(std::acos(a))); };
        return ogn::compute::tryCompute(db.inputs.a().template get<T>(), db.outputs.result().template get<T>(), functor, count);
    }
    else if (operation == db.tokens.ARCTAN) // Arctangent
    {
        auto functor = [](auto const& a, auto& result)
        { result = static_cast<T>(pxr::GfRadiansToDegrees(std::atan(a))); };
        return ogn::compute::tryCompute(db.inputs.a().template get<T>(), db.outputs.result().template get<T>(), functor, count);
    }
    else if (operation == db.tokens.DEGREES) // Degrees
    {
        auto functor = [](auto const& a, auto& result) { result = static_cast<T>(pxr::GfRadiansToDegrees(a)); };
        return ogn::compute::tryCompute(db.inputs.a().template get<T>(), db.outputs.result().template get<T>(), functor, count);
    }
    else if (operation == db.tokens.RADIANS) // Radians
    {
        auto functor = [](auto const& a, auto& result) { result = static_cast<T>(pxr::GfDegreesToRadians(a)); };
        return ogn::compute::tryCompute(db.inputs.a().template get<T>(), db.outputs.result().template get<T>(), functor, count);
    }

    throw ogn::compute::InputError("Operation not one of sin, cos, tan, asin, acos, degrees, radians");
}

template <>
bool tryComputeAssumingType<pxr::GfHalf>(OgnTrigDatabase& db, NameToken operation, size_t count)
{
    if (operation == db.tokens.SIN) // Sine
    {
        auto functor = [](auto const& a, auto& result)
        { result = static_cast<pxr::GfHalf>(static_cast<float>(std::sin(pxr::GfDegreesToRadians(a)))); };
        return ogn::compute::tryCompute(
            db.inputs.a().template get<pxr::GfHalf>(), db.outputs.result().template get<pxr::GfHalf>(), functor, count);
    }
    else if (operation == db.tokens.COS) // Cosine
    {
        auto functor = [](auto const& a, auto& result)
        { result = static_cast<pxr::GfHalf>(static_cast<float>(std::cos(pxr::GfDegreesToRadians(a)))); };
        return ogn::compute::tryCompute(
            db.inputs.a().template get<pxr::GfHalf>(), db.outputs.result().template get<pxr::GfHalf>(), functor, count);
    }
    else if (operation == db.tokens.TAN) // Tangent
    {
        auto functor = [](auto const& a, auto& result)
        { result = static_cast<pxr::GfHalf>(static_cast<float>(std::tan(pxr::GfDegreesToRadians(a)))); };
        return ogn::compute::tryCompute(
            db.inputs.a().template get<pxr::GfHalf>(), db.outputs.result().template get<pxr::GfHalf>(), functor, count);
    }
    else if (operation == db.tokens.ARCSIN) // Arcsine
    {
        auto functor = [](auto const& a, auto& result)
        { result = static_cast<pxr::GfHalf>(static_cast<float>(std::asin(pxr::GfDegreesToRadians(a)))); };
        return ogn::compute::tryCompute(
            db.inputs.a().template get<pxr::GfHalf>(), db.outputs.result().template get<pxr::GfHalf>(), functor, count);
    }
    else if (operation == db.tokens.ARCCOS) // Arccosine
    {
        auto functor = [](auto const& a, auto& result)
        { result = static_cast<pxr::GfHalf>(static_cast<float>(std::acos(pxr::GfDegreesToRadians(a)))); };
        return ogn::compute::tryCompute(
            db.inputs.a().template get<pxr::GfHalf>(), db.outputs.result().template get<pxr::GfHalf>(), functor, count);
    }
    else if (operation == db.tokens.ARCTAN) // Arctangent
    {
        auto functor = [](auto const& a, auto& result)
        { result = static_cast<pxr::GfHalf>(static_cast<float>(std::atan(pxr::GfDegreesToRadians(a)))); };
        return ogn::compute::tryCompute(
            db.inputs.a().template get<pxr::GfHalf>(), db.outputs.result().template get<pxr::GfHalf>(), functor, count);
    }
    else if (operation == db.tokens.DEGREES) // Degrees
    {
        auto functor = [](auto const& a, auto& result)
        { result = static_cast<pxr::GfHalf>(static_cast<float>(pxr::GfRadiansToDegrees(a))); };
        return ogn::compute::tryCompute(
            db.inputs.a().template get<pxr::GfHalf>(), db.outputs.result().template get<pxr::GfHalf>(), functor, count);
    }
    else if (operation == db.tokens.RADIANS) // Radians
    {
        auto functor = [](auto const& a, auto& result)
        { result = static_cast<pxr::GfHalf>(static_cast<float>(pxr::GfDegreesToRadians(a))); };
        return ogn::compute::tryCompute(
            db.inputs.a().template get<pxr::GfHalf>(), db.outputs.result().template get<pxr::GfHalf>(), functor, count);
    }

    throw ogn::compute::InputError("Operation not one of sin, cos, tan, asin, acos, degrees, radians");
}
} // namespace


class OgnTrig

{
public:
    static size_t computeVectorized(OgnTrigDatabase& db, size_t count)
    {
        NameToken const& operation = db.inputs.operation();

        try
        {
            if (tryComputeAssumingType<double>(db, operation, count)) return count;
            else if (tryComputeAssumingType<pxr::GfHalf>(db, operation, count)) return count;
            else if (tryComputeAssumingType<float>(db, operation, count)) return count;
            else
            {
                db.logWarning("Failed to resolve input types");
            }
        }
        catch (std::exception &error)
        {
            db.logError("Operation %s could not be performed : %s", db.tokenToString(operation), error.what());
        }
        return 0;
    }

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

        auto aType = a.iAttribute->getResolvedType(a);

        // Require inputs to be resolved before determining output's type
        if (aType.baseType != BaseDataType::eUnknown)
        {
            std::array<AttributeObj, 2> attrs { a, result };
            node.iNode->resolveCoupledAttributes(node, attrs.data(), attrs.size());

        }
    }
};

REGISTER_OGN_NODE()

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