// Copyright (c) 2021-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 <OgnTransformVectorDatabase.h>
#include <omni/graph/core/ogn/ComputeHelpers.h>
#include <omni/math/linalg/matrix.h>
#include <omni/math/linalg/vec.h>

using omni::math::linalg::matrix4d;
using omni::math::linalg::matrix3d;
using omni::math::linalg::vec3;
using ogn::compute::tryComputeWithArrayBroadcasting;

namespace omni
{
namespace graph
{
namespace nodes
{

namespace
{
    
template<typename T> 
bool tryComputeWithMatrix3(OgnTransformVectorDatabase& db)
{
    auto functor = [&](auto& matrix, auto& vector, auto& result)
    {
        auto& transformMat = *reinterpret_cast<const matrix3d*>(matrix);
        auto sourceVec = vec3<double>(*reinterpret_cast<const vec3<T>*>(vector));
        auto& resultVec = *reinterpret_cast<vec3<T>*>(result);
        // left multiplication by row vector
        resultVec = vec3<T>(sourceVec * transformMat);
    };
    return tryComputeWithArrayBroadcasting<double[9], T[3], T[3]>(db.inputs.matrix(), db.inputs.vector(), db.outputs.result(), functor);

}
template<typename T> 
bool tryComputeWithMatrix4(OgnTransformVectorDatabase& db)
{
    auto functor = [&](auto& matrix, auto& vector, auto& result)
    {
        auto& transformMat = *reinterpret_cast<const matrix4d*>(matrix);
        auto sourceVec = vec3<double>(*reinterpret_cast<const vec3<T>*>(vector));
        auto& resultVec = *reinterpret_cast<vec3<T>*>(result);
        resultVec = vec3<T>(transformMat.Transform(sourceVec));
    };
    return tryComputeWithArrayBroadcasting<double[16], T[3], T[3]>(db.inputs.matrix(), db.inputs.vector(), db.outputs.result(), functor);
}

/* FIXME: GfHalf has no explicit conversion from double to half, so we need to convert to a float first */
template<>
bool tryComputeWithMatrix3<pxr::GfHalf>(OgnTransformVectorDatabase& db)
{
    auto functor = [&](auto& matrix, auto& vector, auto& result)
    {
        auto& transformMat = *reinterpret_cast<const matrix3d*>(matrix);
        auto sourceVec = vec3<double>(*reinterpret_cast<const vec3<pxr::GfHalf>*>(vector));
        auto& resultVec = *reinterpret_cast<vec3<pxr::GfHalf>*>(result);
        // left multiplication by row vector
        resultVec = vec3<pxr::GfHalf>(vec3<float>(sourceVec * transformMat));
    };
    return tryComputeWithArrayBroadcasting<double[9], pxr::GfHalf[3], pxr::GfHalf[3]>(db.inputs.matrix(), db.inputs.vector(), db.outputs.result(), functor);

}
template<> 
bool tryComputeWithMatrix4<pxr::GfHalf>(OgnTransformVectorDatabase& db)
{
    auto functor = [&](auto& matrix, auto& vector, auto& result)
    {
        auto& transformMat = *reinterpret_cast<const matrix4d*>(matrix);
        auto sourceVec = vec3<double>(*reinterpret_cast<const vec3<pxr::GfHalf>*>(vector));
        auto& resultVec = *reinterpret_cast<vec3<pxr::GfHalf>*>(result);
        resultVec =  vec3<pxr::GfHalf>(vec3<float>(transformMat.Transform(sourceVec)));
    };
    return tryComputeWithArrayBroadcasting<double[16], pxr::GfHalf[3], pxr::GfHalf[3]>(db.inputs.matrix(), db.inputs.vector(), db.outputs.result(), functor);
}
} // namespace

class OgnTransformVector
{
public:
    static bool compute(OgnTransformVectorDatabase& db)
    {
        try 
        {
            // matrix3
            if (tryComputeWithMatrix3<double>(db)) return true;
            else if (tryComputeWithMatrix3<float>(db)) return true;
            else if (tryComputeWithMatrix3<pxr::GfHalf>(db)) return true;

            // matrix4
            else if (tryComputeWithMatrix4<double>(db)) return true;
            else if (tryComputeWithMatrix4<float>(db)) return true;
            else if (tryComputeWithMatrix4<pxr::GfHalf>(db)) return true;
            else
            {
                db.logWarning("OgnTransformVector: Failed to resolve input types");
            }
        }
        catch (ogn::compute::InputError &error)
        {
            db.logWarning("OgnTransformVector: %s", error.what());
        }
        return false;
    }
    
    static void onConnectionTypeResolve(const NodeObj& node){
        auto vector = node.iNode->getAttributeByToken(node, inputs::vector.token());
        auto matrix = node.iNode->getAttributeByToken(node, inputs::matrix.token());
        auto result = node.iNode->getAttributeByToken(node, outputs::result.token());

        auto vectorType = vector.iAttribute->getResolvedType(vector);
        auto matrixType = matrix.iAttribute->getResolvedType(matrix);

        // Require vector, matrix to be resolved before determining result's type
        if (vectorType.baseType != BaseDataType::eUnknown && matrixType.baseType != BaseDataType::eUnknown)
        {   
            Type resultType(vectorType.baseType, vectorType.componentCount, std::max(vectorType.arrayDepth, matrixType.arrayDepth), vectorType.role);
            result.iAttribute->setResolvedType(result, resultType);
        }
    }
};

REGISTER_OGN_NODE()

}
}
}
