// 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 <OgnMatrixMultiplyDatabase.h>
#include <omni/graph/core/ogn/ComputeHelpers.h>
#include <carb/logging/Log.h>
#include <omni/math/linalg/matrix.h>
#include <omni/math/linalg/vec.h>

#include <type_traits>

using omni::math::linalg::base_matrix;
using omni::math::linalg::base_vec;
using ogn::compute::tryComputeWithArrayBroadcasting;

namespace omni
{
namespace graph
{
namespace nodes
{

namespace 
{
// Base type T, dimension N (eg: 2,3,4)
template<typename T, size_t N>
bool tryComputeAssumingType(OgnMatrixMultiplyDatabase& db)
{
    auto matrixMatrixFn = [](auto const& a, auto const& b, auto& result)
    {
        const auto& aMatrix = *reinterpret_cast<const base_matrix<T, N>*>(a);
        const auto& bMatrix = *reinterpret_cast<const base_matrix<T, N>*>(b);
        auto& resultMatrix  = *reinterpret_cast<base_matrix<T, N>*>(result);
        resultMatrix = aMatrix * bMatrix;
    };
    auto matrixVectorFn = [](auto const& a, auto const& b, auto& result)
    {
        const auto& aMatrix = *reinterpret_cast<const base_matrix<T, N>*>(a);
        const auto& bVec    = *reinterpret_cast<const base_vec<T, N>*>(b);
        auto& resultVec     = *reinterpret_cast<base_vec<T, N>*>(result);
        resultVec = aMatrix * bVec;
    };
    auto vectorMatrixFn = [](auto const& a, auto const& b, auto& result)
    {
        const auto& aVec    = *reinterpret_cast<const base_vec<T, N>*>(a);
        const auto& bMatrix = *reinterpret_cast<const base_matrix<T, N>*>(b);
        auto& resultVec     = *reinterpret_cast<base_vec<T, N>*>(result);
        resultVec = aVec * bMatrix;
    };
    auto vectorVectorFn = [](auto const& a, auto const& b, auto& result)
    {
        const auto& aVec = *reinterpret_cast<const base_vec<T, N>*>(a);
        const auto& bVec = *reinterpret_cast<const base_vec<T, N>*>(b);
        result = aVec * bVec;
    };
    if (N >= 3 && tryComputeWithArrayBroadcasting<T[N*N]>(db.inputs.a(), db.inputs.b(), db.outputs.output(), matrixMatrixFn)) 
        return true;
    else if (N >= 3 && tryComputeWithArrayBroadcasting<T[N*N], T[N], T[N]>(db.inputs.a(), db.inputs.b(), db.outputs.output(), matrixVectorFn)) 
        return true;
    else if (N >= 3 && tryComputeWithArrayBroadcasting<T[N], T[N*N], T[N]>(db.inputs.a(), db.inputs.b(), db.outputs.output(), vectorMatrixFn)) 
        return true;
    else if (tryComputeWithArrayBroadcasting<T[N], T[N], T>(db.inputs.a(), db.inputs.b(), db.outputs.output(), vectorVectorFn)) 
        return true;
    return false;
}
} // namespace

class OgnMatrixMultiply
{
public:
    static bool compute(OgnMatrixMultiplyDatabase& db)
    {
        try 
        {
            if (tryComputeAssumingType<double, 2>(db)) return true;
            else if (tryComputeAssumingType<double, 3>(db)) return true;
            else if (tryComputeAssumingType<double, 4>(db)) return true;
            else if (tryComputeAssumingType<float, 2>(db)) return true;
            else if (tryComputeAssumingType<float, 3>(db)) return true;
            else if (tryComputeAssumingType<float, 4>(db)) return true;
            else if (tryComputeAssumingType<pxr::GfHalf, 2>(db)) return true; 
            else if (tryComputeAssumingType<pxr::GfHalf, 3>(db)) return true; 
            else if (tryComputeAssumingType<pxr::GfHalf, 4>(db)) return true; 
            else
            {
                db.logWarning("OgnMatrixMultiply: Failed to resolve input types");
            }
        }
        catch (ogn::compute::InputError &error)
        {
            db.logWarning("OgnMatrixMultiply: %s", error.what());
        }
        return false;
    }

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

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

        // Require inputs to be resolved before determining output's type
        if (aType.baseType != BaseDataType::eUnknown && bType.baseType != BaseDataType::eUnknown)
        {
            auto biggerDim = std::max(aType.componentCount, bType.componentCount);
            auto smallerDim = std::min(aType.componentCount, bType.componentCount);

            if (biggerDim != smallerDim && biggerDim != smallerDim * smallerDim)
            {
                CARB_LOG_WARN_ONCE("OgnMatrixMultiply: Inputs are not compatible with tuple counts %d and %d", aType.componentCount, bType.componentCount);
                return;
            }
            
            // Vector4 * Matrix4 = Vector4, Matrix4 * Vector4 = Vector4 and etc.
            auto tupleCount = smallerDim;
            auto role = aType.componentCount < bType.componentCount ? aType.role : bType.role;
            if (smallerDim == biggerDim && smallerDim <= 4)
            {
                // equivalent to dot product of two vectors
                tupleCount = 1;
                role = AttributeRole::eNone;
            }
            std::array<AttributeObj, 3> attrs { a, b, output };
            std::array<uint8_t, 3> tupleCounts {
                aType.componentCount,
                bType.componentCount,
                tupleCount
            };
            std::array<uint8_t, 3> arrayDepths {
                aType.arrayDepth,
                bType.arrayDepth,
                std::max(aType.arrayDepth, bType.arrayDepth)
            };
            std::array<AttributeRole, 3> rolesBuf {
                aType.role,
                bType.role,
                role
            };
            node.iNode->resolvePartiallyCoupledAttributes(node, attrs.data(), tupleCounts.data(), 
                                                        arrayDepths.data(), rolesBuf.data(), attrs.size());
        }
    }
};

REGISTER_OGN_NODE()

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