// 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 <OgnGetMatrix4TranslationDatabase.h>
#include <omni/math/linalg/vec.h>
#include <omni/math/linalg/matrix.h>

using omni::math::linalg::matrix4d;
using omni::math::linalg::vec3d;

namespace omni
{
namespace graph
{
namespace nodes
{

namespace
{
    void getTranslation(const double input[16], double result[4])
    {
        vec3d translation = reinterpret_cast<const matrix4d*>(input)->ExtractTranslation();
        memcpy(result, translation.data(), sizeof(double) * 3);
    }
}

class OgnGetMatrix4Translation
{
public:
    static bool computeVectorized(OgnGetMatrix4TranslationDatabase& db, size_t count)
    {
        const auto& matrixInput = db.inputs.matrix();

        bool inOk = false;
        bool failed = true;

        // Singular case
        if (matrixInput.type().arrayDepth == 0)
        {
            if (auto matrix = matrixInput.get<double[16]>())
            {
                inOk = true;
                auto translationOutput = db.outputs.translation();
                if (auto output = translationOutput.get<double[3]>())
                {
                    auto matrices = matrix.vectorized(count);
                    auto outputs = output.vectorized(count);
                    for (size_t i = 0; i < count; ++i)
                        getTranslation(matrices[i], outputs[i]);
                    failed = false;
                }
            }
        }
        // Array case
        else
        {
            for (size_t inst = 0; inst < count; ++inst)
            {
                if (auto matrices = db.inputs.matrix(inst).get<double[][16]>())
                {
                    inOk = true;
                    if (auto output = db.outputs.translation(inst).get<double[][3]>())
                    {
                        output->resize(matrices->size());
                        for (size_t i = 0; i < matrices.size(); i++)
                        {
                            getTranslation((*matrices)[i], (*output)[i]);
                        }
                        failed = false;
                    }
                }
            }
        }


        if (!inOk)
        {
            db.logError("Input type for matrix input not supported");
            return false;
        }
        
        if (failed)
        {
            db.logError("Input and output depths need to align: Matrix input with depth %s, output with depth %s", 
                    matrixInput.type().arrayDepth, db.outputs.translation().type().arrayDepth);
            return false;
        }

        return true;
    }
    
    static void onConnectionTypeResolve(const NodeObj& node) {
        auto start = node.iNode->getAttribute(node, OgnGetMatrix4TranslationAttributes::inputs::matrix.m_name);
        auto result = node.iNode->getAttribute(node, OgnGetMatrix4TranslationAttributes::outputs::translation.m_name);
        auto startType = start.iAttribute->getResolvedType(start);

        std::array<AttributeObj, 2> attrs { start, result };
        std::array<uint8_t, 2> tupleCounts { 16, 3 };
        std::array<uint8_t, 2> arrayDepths { startType.arrayDepth, startType.arrayDepth };
        std::array<AttributeRole, 2> rolesBuf { AttributeRole::eNone, AttributeRole::eVector };
        node.iNode->resolvePartiallyCoupledAttributes(node, attrs.data(), tupleCounts.data(), 
                                                    arrayDepths.data(), rolesBuf.data(), attrs.size());
    }
};

REGISTER_OGN_NODE()

}
}
}
